/*
 * (C) Copyright 2013 ECMWF.
 *
 * This software is licensed under the terms of the Apache Licence Version 2.0
 * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 * In applying this licence, ECMWF does not waive the privileges and immunities
 * granted to it by virtue of its status as an intergovernmental organisation
 * nor does it submit to any jurisdiction.
 */

#include <cmath>
#include <iomanip>
#include <iostream>

#include "eckit/config/Resource.h"

#include "atlas/array/MakeView.h"
#include "atlas/field/Field.h"
#include "atlas/field/FieldSet.h"
#include "atlas/grid/Distribution.h"
#include "atlas/grid/Grid.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/meshgenerator.h"
#include "atlas/numerics/Nabla.h"
#include "atlas/numerics/fvm/Method.h"
#include "atlas/option.h"
#include "atlas/output/Gmsh.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/util/Config.h"
#include "atlas/util/Constants.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/Earth.h"

#include "tests/AtlasTestEnvironment.h"

using namespace eckit;
using namespace atlas::numerics;
using namespace atlas::meshgenerator;
using namespace atlas::grid;

namespace atlas {
namespace test {


int test_levels() {
    static int levels = []() {
        int levels = eckit::Resource<int>( "--levels", 0 );
        Log::info() << "levels = " << levels << std::endl;
        return levels;
    }();
    return levels;
}

static std::string griduid() {
    return "Slat20";
}

array::LocalView<double, 3> make_vectorview( Field& field ) {
    using array::Range;
    return field.levels() ? array::make_view<double, 3>( field ).slice( Range::all(), Range::all(), Range::all() )
                          : array::make_view<double, 2>( field ).slice( Range::all(), Range::dummy(), Range::all() );
}

array::LocalView<double, 2> make_scalarview( Field& field ) {
    using array::Range;
    return field.levels() ? array::make_view<double, 2>( field ).slice( Range::all(), Range::all() )
                          : array::make_view<double, 1>( field ).slice( Range::all(), Range::dummy() );
}


#define print_min_max_mean( name )                                                \
    do {                                                                          \
        Log::info() << #name << std::endl;                                        \
        Log::info() << std::setprecision( 18 ) << "  min  " << min << std::endl;  \
        Log::info() << std::setprecision( 18 ) << "  max  " << max << std::endl;  \
        Log::info() << std::setprecision( 18 ) << "  mean " << mean << std::endl; \
    } while ( 0 )


//-----------------------------------------------------------------------------

double dual_volume( const Mesh& mesh ) {
    const mesh::Nodes& nodes = mesh.nodes();
    int nb_nodes             = nodes.size();
    auto dual_volumes        = array::make_view<double, 1>( nodes.field( "dual_volumes" ) );
    auto is_ghost            = array::make_view<int, 1>( nodes.ghost() );
    double area              = 0;
    for ( int node = 0; node < nb_nodes; ++node ) {
        if ( !is_ghost( node ) ) {
            area += dual_volumes( node );
        }
    }

    mpi::comm().allReduceInPlace( area, eckit::mpi::sum() );

    return area;
}

/// @brief Compute magnitude of flow with rotation-angle beta
/// (beta=0 --> zonal, beta=pi/2 --> meridional)
void rotated_flow( const fvm::Method& fvm, Field& field, const double& beta ) {
    const double radius  = fvm.radius();
    const double USCAL   = 20.;
    const double pvel    = USCAL / radius;
    const double deg2rad = M_PI / 180.;

    auto lonlat_deg = array::make_view<double, 2>( fvm.mesh().nodes().lonlat() );
    auto var        = make_vectorview( field );

    const idx_t nnodes = fvm.mesh().nodes().size();
    const idx_t nlev   = var.shape( 1 );
    for ( idx_t jnode = 0; jnode < nnodes; ++jnode ) {
        double x = lonlat_deg( jnode, LON ) * deg2rad;
        double y = lonlat_deg( jnode, LAT ) * deg2rad;
        double Ux =
            pvel * ( std::cos( beta ) + std::tan( y ) * std::cos( x ) * std::sin( beta ) ) * radius * std::cos( y );
        double Uy = -pvel * std::sin( x ) * std::sin( beta ) * radius;
        for ( idx_t jlev = 0; jlev < nlev; ++jlev ) {
            var( jnode, jlev, LON ) = Ux;
            var( jnode, jlev, LAT ) = Uy;
        }
    }
}

/// @brief Compute magnitude of flow with rotation-angle beta
/// (beta=0 --> zonal, beta=pi/2 --> meridional)
void rotated_flow_magnitude( const fvm::Method& fvm, Field& field, const double& beta ) {
    const double radius  = fvm.radius();
    const double USCAL   = 20.;
    const double pvel    = USCAL / radius;
    const double deg2rad = M_PI / 180.;

    auto lonlat_deg = array::make_view<double, 2>( fvm.mesh().nodes().lonlat() );
    auto var        = make_scalarview( field );

    const idx_t nnodes = fvm.mesh().nodes().size();
    const idx_t nlev   = var.shape( 1 );
    for ( idx_t jnode = 0; jnode < nnodes; ++jnode ) {
        double x = lonlat_deg( jnode, LON ) * deg2rad;
        double y = lonlat_deg( jnode, LAT ) * deg2rad;
        double Ux =
            pvel * ( std::cos( beta ) + std::tan( y ) * std::cos( x ) * std::sin( beta ) ) * radius * std::cos( y );
        double Uy = -pvel * std::sin( x ) * std::sin( beta ) * radius;
        for ( idx_t jlev = 0; jlev < nlev; ++jlev ) {
            var( jnode, jlev ) = std::sqrt( Ux * Ux + Uy * Uy );
        }
    }
}

//-----------------------------------------------------------------------------

CASE( "test_factory" ) {
    EXPECT( NablaFactory::has( "fvm" ) );
}

CASE( "test_build" ) {
    Log::info() << "test_build" << std::endl;
    MeshGenerator meshgenerator( "structured" );
    Mesh mesh      = meshgenerator.generate( Grid( "O16" ) );
    const double R = util::Earth::radius();
    fvm::Method fvm( mesh, util::Config( "radius", R ) );
    Nabla nabla( fvm );

    double spherical_area = 360. * 180.;
    EXPECT( eckit::types::is_approximately_equal( dual_volume( mesh ), spherical_area, 5.0 ) );
}

CASE( "test_grad" ) {
    Log::info() << "test_grad" << std::endl;
    auto radius = option::radius( "Earth" );
    Grid grid( griduid() );
    MeshGenerator meshgenerator( "structured" );
    Mesh mesh = meshgenerator.generate( grid, Distribution( grid, Partitioner( "equal_regions" ) ) );
    fvm::Method fvm( mesh, radius | option::levels( test_levels() ) );
    Nabla nabla( fvm );

    idx_t nnodes = mesh.nodes().size();
    idx_t nlev   = std::max( 1, test_levels() );

    FieldSet fields;
    fields.add( fvm.node_columns().createField<double>( option::name( "scalar" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "rscalar" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "grad" ) | option::variables( 2 ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "rgrad" ) | option::variables( 2 ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "xder" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "yder" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "rxder" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "ryder" ) ) );

    EXPECT( fields["scalar"].rank() == ( test_levels() ? 2 : 1 ) );
    EXPECT( fields["grad"].rank() == fields["scalar"].rank() + 1 );
    EXPECT( fields["scalar"].levels() == test_levels() );
    EXPECT( fields["grad"].levels() == test_levels() );

    rotated_flow_magnitude( fvm, fields["scalar"], 0. );
    rotated_flow_magnitude( fvm, fields["rscalar"], M_PI_2 * 0.75 );

    nabla.gradient( fields["scalar"], fields["grad"] );
    nabla.gradient( fields["rscalar"], fields["rgrad"] );
    auto xder        = make_scalarview( fields["xder"] );
    auto yder        = make_scalarview( fields["yder"] );
    auto rxder       = make_scalarview( fields["rxder"] );
    auto ryder       = make_scalarview( fields["ryder"] );
    const auto grad  = make_vectorview( fields["grad"] );
    const auto rgrad = make_vectorview( fields["rgrad"] );
    for ( idx_t jnode = 0; jnode < nnodes; ++jnode ) {
        for ( idx_t jlev = 0; jlev < nlev; ++jlev ) {
            xder( jnode, jlev )  = grad( jnode, jlev, LON );
            yder( jnode, jlev )  = grad( jnode, jlev, LAT );
            rxder( jnode, jlev ) = rgrad( jnode, jlev, LON );
            ryder( jnode, jlev ) = rgrad( jnode, jlev, LAT );
        }
    }

    // output to gmsh
    {
        fvm.node_columns().haloExchange( fields );
        output::Gmsh( grid.name() + ".msh" ).write( mesh );
        output::Gmsh gmsh_fields( grid.name() + "_fields.msh" );
        gmsh_fields.write( fields["scalar"] );
        gmsh_fields.write( fields["xder"] );
        gmsh_fields.write( fields["yder"] );
        gmsh_fields.write( fields["rscalar"] );
        gmsh_fields.write( fields["rxder"] );
        gmsh_fields.write( fields["ryder"] );
    }

    double min, max, mean;
    idx_t N;

    fvm.node_columns().minimum( fields["xder"], min );
    fvm.node_columns().maximum( fields["xder"], max );
    fvm.node_columns().mean( fields["xder"], mean, N );
    print_min_max_mean( "xder" );
    EXPECT_APPROX_EQ( min, 0., 1.e-20 );
    EXPECT_APPROX_EQ( max, 0., 1.e-20 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );

    fvm.node_columns().minimum( fields["yder"], min );
    fvm.node_columns().maximum( fields["yder"], max );
    fvm.node_columns().mean( fields["yder"], mean, N );
    print_min_max_mean( "yder" );
    EXPECT_APPROX_EQ( min, -3.1141489788326316614e-06 );
    EXPECT_APPROX_EQ( max, 3.1141489788326316614e-06 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );


    fvm.node_columns().minimum( fields["rxder"], min );
    fvm.node_columns().maximum( fields["rxder"], max );
    fvm.node_columns().mean( fields["rxder"], mean, N );
    print_min_max_mean( "rxder" );
    EXPECT_APPROX_EQ( min, -3.02863817262107e-06 );
    EXPECT_APPROX_EQ( max, +3.02863817262107e-06 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );

    fvm.node_columns().minimum( fields["ryder"], min );
    fvm.node_columns().maximum( fields["ryder"], max );
    fvm.node_columns().mean( fields["ryder"], mean, N );
    print_min_max_mean( "ryder" );
    EXPECT_APPROX_EQ( min, -3.114148978832633e-06 );
    EXPECT_APPROX_EQ( max, +3.114148978832633e-06 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );
}

CASE( "test_div" ) {
    Log::info() << "test_div" << std::endl;
    const double radius = util::Earth::radius();
    //  const double radius = 1.;
    Grid grid( griduid() );
    MeshGenerator meshgenerator( "structured" );
    Mesh mesh = meshgenerator.generate( grid, Distribution( grid, Partitioner( "equal_regions" ) ) );
    fvm::Method fvm( mesh, util::Config( "radius", radius ) | option::levels( test_levels() ) );
    Nabla nabla( fvm );

    FieldSet fields;
    fields.add( fvm.node_columns().createField<double>( option::name( "wind" ) | option::variables( 2 ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "div" ) ) );

    rotated_flow( fvm, fields["wind"], M_PI_2 * 0.75 );

    nabla.divergence( fields["wind"], fields["div"] );

    // output to gmsh
    {
        fvm.node_columns().haloExchange( fields );
        output::Gmsh gmsh( grid.name() + "_fields.msh", "a" );
        gmsh.write( fields["wind"] );
        gmsh.write( fields["div"] );
    }

    double min, max, mean;
    idx_t N;
    fvm.node_columns().minimum( fields["div"], min );
    fvm.node_columns().maximum( fields["div"], max );
    fvm.node_columns().mean( fields["div"], mean, N );

    // Divergence free flow!
    print_min_max_mean( "div" );
    EXPECT_APPROX_EQ( min, 0., 1.e-18 );
    EXPECT_APPROX_EQ( max, 0., 1.e-18 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );
}

CASE( "test_curl" ) {
    Log::info() << "test_curl" << std::endl;
    const double radius = util::Earth::radius();
    //  const double radius = 1.;
    Grid grid( griduid() );
    MeshGenerator meshgenerator( "structured" );
    Mesh mesh = meshgenerator.generate( grid, Distribution( grid, Partitioner( "equal_regions" ) ) );
    fvm::Method fvm( mesh, util::Config( "radius", radius ) | option::levels( test_levels() ) );
    Nabla nabla( fvm );

    FieldSet fields;
    fields.add( fvm.node_columns().createField<double>( option::name( "wind" ) | option::variables( 2 ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "vor" ) ) );

    rotated_flow( fvm, fields["wind"], M_PI_2 * 0.75 );

    nabla.curl( fields["wind"], fields["vor"] );

    fields.add( fvm.node_columns().createField<double>( option::name( "windgrad" ) | option::variables( 2 * 2 ) ) );
    nabla.gradient( fields["wind"], fields["windgrad"] );

    fields.add( fvm.node_columns().createField<double>( option::name( "windX" ) | option::levels( false ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "windY" ) | option::levels( false ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "windXgradX" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "windXgradY" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "windYgradX" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "windYgradY" ) ) );
    auto wind     = make_vectorview( fields["wind"] );
    auto windgrad = make_vectorview( fields["windgrad"] );

    auto windX      = array::make_view<double, 1>( fields["windX"] );
    auto windY      = array::make_view<double, 1>( fields["windY"] );
    auto windXgradX = make_scalarview( fields["windXgradX"] );
    auto windXgradY = make_scalarview( fields["windXgradY"] );
    auto windYgradX = make_scalarview( fields["windYgradX"] );
    auto windYgradY = make_scalarview( fields["windYgradY"] );
    for ( idx_t j = 0; j < windX.size(); ++j ) {
        static const idx_t lev0 = 0;
        static const idx_t XdX  = XX * 2 + XX;
        static const idx_t XdY  = XX * 2 + YY;
        static const idx_t YdX  = YY * 2 + XX;
        static const idx_t YdY  = YY * 2 + YY;
        windX( j )              = wind( j, lev0, XX );
        windY( j )              = wind( j, lev0, YY );
        windXgradX( j, lev0 )   = windgrad( j, lev0, XdX );
        windXgradY( j, lev0 )   = windgrad( j, lev0, XdY );
        windYgradX( j, lev0 )   = windgrad( j, lev0, YdX );
        windYgradY( j, lev0 )   = windgrad( j, lev0, YdY );
    }

    // output to gmsh
    {
        fvm.node_columns().haloExchange( fields );
        output::Gmsh gmsh( grid.name() + "_fields.msh", "a" );
        gmsh.write( fields["vor"] );
        gmsh.write( fields["windX"] );
        gmsh.write( fields["windXgradX"] );
        gmsh.write( fields["windXgradY"] );
        gmsh.write( fields["windY"] );
        gmsh.write( fields["windYgradX"] );
        gmsh.write( fields["windYgradY"] );
        gmsh.write( fields["windgrad"] );
    }

    double min, max, mean;
    idx_t N;

    // Vorticity!
    fvm.node_columns().minimum( fields["vor"], min );
    fvm.node_columns().maximum( fields["vor"], max );
    fvm.node_columns().mean( fields["vor"], mean, N );
    print_min_max_mean( "vor" );
    EXPECT_APPROX_EQ( min, -6.257451225821150e-06 );
    EXPECT_APPROX_EQ( max, 6.257451225821150e-06 );
    EXPECT_APPROX_EQ( mean, 0., 1.e-20 );
}

CASE( "test_lapl" ) {
    Log::info() << "test_lapl" << std::endl;
    const double radius = util::Earth::radius();
    //  const double radius = 1.;
    Grid grid( griduid() );
    MeshGenerator meshgenerator( "structured" );
    Mesh mesh = meshgenerator.generate( grid, Distribution( grid, Partitioner( "equal_regions" ) ) );
    fvm::Method fvm( mesh, util::Config( "radius", radius ) | option::levels( test_levels() ) );
    Nabla nabla( fvm );

    FieldSet fields;
    fields.add( fvm.node_columns().createField<double>( option::name( "scal" ) ) );
    fields.add( fvm.node_columns().createField<double>( option::name( "lapl" ) ) );

    rotated_flow_magnitude( fvm, fields["scal"], M_PI_2 * 0.75 );

    nabla.laplacian( fields["scal"], fields["lapl"] );

    // output to gmsh
    {
        fvm.node_columns().haloExchange( fields );
        output::Gmsh gmsh( grid.name() + "_fields.msh", "a" );
        gmsh.write( fields["lapl"] );
    }

    double min, max, mean;
    idx_t N;
    fvm.node_columns().minimum( fields["lapl"], min );
    fvm.node_columns().maximum( fields["lapl"], max );
    fvm.node_columns().mean( fields["lapl"], mean, N );
    print_min_max_mean( "lapl" );
    EXPECT_APPROX_EQ( min, -6.4088005677811607095e-13 );
    EXPECT_APPROX_EQ( max, 9.8984499569639476135e-12 );
    EXPECT_APPROX_EQ( mean, -1.03409e-13 );
}

//-----------------------------------------------------------------------------

}  // namespace test
}  // namespace atlas

int main( int argc, char** argv ) {
    return atlas::test::run( argc, argv );
}
