/*
 * (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 <algorithm>
#include <iomanip>

#include "eckit/linalg/LinearAlgebra.h"
#include "eckit/linalg/Vector.h"
#include "eckit/types/Types.h"

#include "atlas/array.h"
#include "atlas/field/Field.h"
#include "atlas/functionspace/NodeColumns.h"
#include "atlas/functionspace/PointCloud.h"
#include "atlas/functionspace/StructuredColumns.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/grid/StructuredGrid.h"
#include "atlas/interpolation.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/meshgenerator.h"
#include "atlas/output/Gmsh.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/MicroDeg.h"

#include "CubicInterpolationPrototype.h"
#include "tests/AtlasTestEnvironment.h"

using namespace atlas::functionspace;
using namespace atlas::grid;
using namespace atlas::util;
using namespace atlas::array;

namespace atlas {
namespace test {

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

std::vector<double> IFS_full_levels_uniform( idx_t nlev ) {
    std::vector<double> zcoord( nlev );
    double dzcoord = 1. / double( nlev );
    for ( idx_t jlev = 0; jlev < nlev; ++jlev ) {
        zcoord[jlev] = 0.5 * dzcoord + jlev * dzcoord;
    }
    return zcoord;
}

std::vector<double> zrange( idx_t nlev, double min, double max ) {
    std::vector<double> zcoord( nlev );

    double dzcoord = ( max - min ) / double( nlev - 1 );
    for ( idx_t jlev = 0; jlev < nlev; ++jlev ) {
        zcoord[jlev] = min + jlev * dzcoord;
    }
    return zcoord;
}

double cubic( double x, double min, double max ) {
    double x0   = min;
    double x1   = 0.5 * ( max + min );
    double x2   = max;
    double xmax = 0.5 * ( x0 + x1 );
    return ( x - x0 ) * ( x - x1 ) * ( x - x2 ) / ( ( xmax - x0 ) * ( xmax - x1 ) * ( xmax - x2 ) );
}

CASE( "test vertical cubic interpolation" ) {
    idx_t nlev    = 10;
    auto vertical = Vertical{nlev, IFS_full_levels_uniform( nlev )};
    bool limiter  = true;
    CubicVerticalInterpolation interpolate( vertical, util::Config( "limiter", limiter ) );
    std::vector<double> departure_points{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.3246};
    array::ArrayT<double> array( nlev );
    auto view = array::make_view<double, 1>( array );
    Log::info() << "source:" << std::endl;
    for ( idx_t k = 0; k < nlev; ++k ) {
        view( k ) = cubic( vertical( k ), 0., 1. );
        Log::info() << "  " << vertical( k ) << " : " << view( k ) << std::endl;
    }

    Log::info() << "interpolation:" << std::endl;
    for ( auto p : departure_points ) {
        Log::info() << "  " << std::setw( 6 ) << p << " :    " << std::setw( 12 ) << interpolate( p, view )
                    << "     expected : " << cubic( p, 0., 1. ) << std::endl;
        if ( p >= vertical.front() && p <= vertical.back() && !limiter ) {
            EXPECT( eckit::types::is_approximately_equal( interpolate( p, view ), cubic( p, 0., 1. ) ) );
        }
    }
}

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

CASE( "test horizontal cubic interpolation" ) {
    //if ( mpi::comm().size() == 1 ) {
    std::string gridname = eckit::Resource<std::string>( "--grid", "O8" );

    StructuredGrid grid( gridname );
    int halo = eckit::Resource<int>( "--halo", 2 );
    Config config;
    config.set( "halo", halo );
    config.set( "levels", 9 );
    config.set( "periodic_points", true );
    StructuredColumns fs( grid, Partitioner( "equal_regions" ), config );

    Field field = fs.createField<double>( option::levels( 0 ) );
    auto f      = array::make_view<double, 1>( field );
    auto xy     = array::make_view<double, 2>( fs.xy() );

    auto fx  = []( double x ) { return cubic( x, 0., 360. ); };
    auto fy  = []( double y ) { return cubic( y, -90., 90. ); };
    auto fxy = [fx, fy]( double x, double y ) { return fx( x ) * fy( y ); };

    for ( idx_t j = 0; j < fs.size(); ++j ) {
        f( j ) = fxy( xy( j, XX ), xy( j, YY ) );
    }

    CubicHorizontalInterpolation cubic_interpolation( fs );

    auto departure_points = {
        PointXY( 0.13257, 45.6397 ),
    };
    for ( auto p : departure_points ) {
        Log::info() << p << "  -->  " << cubic_interpolation( p.x(), p.y(), f ) << std::endl;
        EXPECT( is_approximately_equal( cubic_interpolation( p.x(), p.y(), f ), fxy( p.x(), p.y() ) ) );
    }
}

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

bool operator==( const eckit::linalg::Triplet& t1, const eckit::linalg::Triplet& t2 ) {
    if ( t1.row() != t2.row() ) {
        return false;
    }
    if ( t1.col() != t2.col() ) {
        return false;
    }
    if ( !is_approximately_equal( t1.value(), t2.value() ) ) {
        return false;
    }

    return true;
}
bool operator!=( const eckit::linalg::Triplet& t1, const eckit::linalg::Triplet& t2 ) {
    return !( t1 == t2 );
}
bool operator==( const std::vector<eckit::linalg::Triplet>& t1, const std::vector<eckit::linalg::Triplet>& t2 ) {
    if ( t1.size() != t2.size() ) {
        return false;
    }
    for ( size_t i = 0; i < t1.size(); ++i ) {
        if ( t1[i] != t2[i] ) {
            return false;
        }
    }
    return true;
}
//std::ostream& operator<<( std::ostream& out, const LocalView<double, 1>& array ) {
//    out << "{ ";
//    for ( idx_t i = 0; i < array.size(); ++i ) {
//        out << array( i );
//        if ( i < array.size() - 1 ) {
//            out << ",";
//        }
//        out << " ";
//    }
//    out << "}";
//    return out;
//}

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

CASE( "test horizontal cubic interpolation triplets" ) {
    //if ( mpi::comm().size() == 1 ) {
    std::string gridname = eckit::Resource<std::string>( "--grid", "O8" );

    StructuredGrid grid( gridname );
    int halo = eckit::Resource<int>( "--halo", 2 );
    Config config;
    config.set( "halo", halo );
    config.set( "levels", 9 );
    config.set( "periodic_points", true );
    StructuredColumns fs( grid, Partitioner( "equal_regions" ), config );

    Field field = fs.createField<double>( option::levels( 0 ) );
    auto f      = array::make_view<double, 1>( field );
    auto xy     = array::make_view<double, 2>( fs.xy() );

    for ( idx_t j = 0; j < fs.size(); ++j ) {
        f( j ) = xy( j, XX );
    }

    CubicHorizontalInterpolation cubic_interpolation( fs );

    auto departure_points = PointCloud{{
        {0.13257, 45.6397},
        {360., -90.},
    }};
    auto departure_lonlat = make_view<double, 2>( departure_points.lonlat() );

    CubicHorizontalInterpolation::WorkSpace ws;
    CubicHorizontalInterpolation::Triplets triplets;
    for ( idx_t row = 0; row < departure_points.size(); ++row ) {
        auto triplets_row =
            cubic_interpolation.compute_triplets( row, departure_lonlat( row, XX ), departure_lonlat( row, YY ), ws );
        Log::info() << departure_lonlat.slice( row, array::Range::all() ) << "  -->  {\n";
        for ( auto triplet : triplets_row ) {
            Log::info() << "    " << triplet << " ,\n";
        }
        Log::info() << " } " << std::endl;
        std::copy( triplets_row.begin(), triplets_row.end(), std::back_inserter( triplets ) );
    }

    auto triplets2 = cubic_interpolation.reserve_triplets( departure_points.size() );
    {
        idx_t row{0};
        for ( auto p : departure_points.iterate().xy() ) {
            cubic_interpolation.insert_triplets( row++, p.x(), p.y(), triplets2, ws );
        }
    }


    EXPECT( triplets2 == triplets );

    eckit::linalg::SparseMatrix matrix( departure_points.size(), fs.size(), triplets2 );
    Log::info() << matrix << std::endl;

    std::vector<double> tgt( departure_points.size() );
    eckit::linalg::Vector v_src( const_cast<double*>( f.data() ), f.size() );
    eckit::linalg::Vector v_tgt( tgt.data(), tgt.size() );
    eckit::linalg::LinearAlgebra::backend().spmv( matrix, v_src, v_tgt );
    Log::info() << "output = " << tgt << std::endl;
}

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

CASE( "test 3d cubic interpolation" ) {
    const double tolerance = 1.e-15;

    //if ( mpi::comm().size() == 1 ) {
    std::string gridname = eckit::Resource<std::string>( "--grid", "O8" );
    idx_t nlev           = 11;

    Vertical vertical( nlev, zrange( nlev, 0., 1. ) );
    Log::info() << zrange( nlev, 0., 1. ) << std::endl;
    ;

    StructuredGrid grid( gridname );
    int halo = eckit::Resource<int>( "--halo", 2 );
    Config config;
    config.set( "halo", halo );
    config.set( "periodic_points", true );
    StructuredColumns fs( grid, vertical, Partitioner( "equal_regions" ), config );

    Field input   = fs.createField<double>();
    auto f        = array::make_view<double, 2>( input );
    auto xy       = array::make_view<double, 2>( fs.xy() );
    const auto& z = fs.vertical();

    auto fx = []( double x ) { return cubic( x, 0., 360. ); };
    auto fy = []( double y ) { return cubic( y, -90., 90. ); };
    auto fz = []( double z ) { return cubic( z, 0., 1. ); };
    auto fp = [fx, fy, fz]( const PointXYZ& p ) { return fx( p.x() ) * fy( p.y() ) * fz( p.z() ); };

    for ( idx_t n = 0; n < fs.size(); ++n ) {
        for ( idx_t k = fs.k_begin(); k < fs.k_end(); ++k ) {
            PointXYZ p{
                xy( n, XX ),
                xy( n, YY ),
                z( k ),
            };
            f( n, k ) = fp( p );
        }
    }
    input.set_dirty( false );  // to avoid halo-exchange


    auto departure_points = PointCloud( {
        {90., -45., 0.25}, {0., -45., 0.25},  {180., -45., 0.25}, {360., -45., 0.25}, {90., -90., 0.25},
        {90., 0., 0.25},   {90., 90., 0.25},  {10., -45., 0.25},  {20., -45., 0.25},  {30., -45., 0.25},
        {40., -45., 0.25}, {50., -45., 0.25}, {60., -45., 0.25},  {70., -45., 0.25},  {80., -45., 0.25},
        {90., -45., 0.25}, {10., -10., 0.25}, {20., -20., 0.25},  {30., -30., 0.25},  {40., -40., 0.25},
        {50., -50., 0.25}, {60., -60., 0.25}, {70., -70., 0.25},  {80., -80., 0.25},  {90., -90., 0.25},
        {60., -60., 0.6},  {90., -45., 0.16}, {90., -45., 0.6},   {90., -45., 1.},    {90., -45., 0.},
        {90., -45., 0.1},  {45., 33., 0.7},   {359., 20., 0.1},   {360., 88., 0.9},   {0.1, -89., 1.},

        {90., -45., 0.25}, {0., -45., 0.25},  {180., -45., 0.25}, {360., -45., 0.25}, {90., -90., 0.25},
        {90., 0., 0.25},   {90., 90., 0.25},  {10., -45., 0.25},  {20., -45., 0.25},  {30., -45., 0.25},
        {40., -45., 0.25}, {50., -45., 0.25}, {60., -45., 0.25},  {70., -45., 0.25},  {80., -45., 0.25},
        {90., -45., 0.25}, {10., -10., 0.25}, {20., -20., 0.25},  {30., -30., 0.25},  {40., -40., 0.25},
        {50., -50., 0.25}, {60., -60., 0.25}, {70., -70., 0.25},  {80., -80., 0.25},  {90., -90., 0.25},
        {60., -60., 0.6},  {90., -45., 0.16}, {90., -45., 0.6},   {90., -45., 1.},    {90., -45., 0.},
        {90., -45., 0.1},  {45., 33., 0.7},   {359., 20., 0.1},   {360., 88., 0.9},   {0.1, -89., 1.},

        {90., -45., 0.25}, {0., -45., 0.25},  {180., -45., 0.25}, {360., -45., 0.25}, {90., -90., 0.25},
        {90., 0., 0.25},   {90., 90., 0.25},  {10., -45., 0.25},  {20., -45., 0.25},  {30., -45., 0.25},
        {40., -45., 0.25}, {50., -45., 0.25}, {60., -45., 0.25},  {70., -45., 0.25},  {80., -45., 0.25},
        {90., -45., 0.25}, {10., -10., 0.25}, {20., -20., 0.25},  {30., -30., 0.25},  {40., -40., 0.25},
        {50., -50., 0.25}, {60., -60., 0.25}, {70., -70., 0.25},  {80., -80., 0.25},  {90., -90., 0.25},
        {60., -60., 0.6},  {90., -45., 0.16}, {90., -45., 0.6},   {90., -45., 1.},    {90., -45., 0.},
        {90., -45., 0.1},  {45., 33., 0.7},   {359., 20., 0.1},   {360., 88., 0.9},   {0.1, -89., 1.},
    } );

    SECTION( "prototype" ) {
        Cubic3DInterpolation cubic_interpolation( fs );

        for ( auto p : departure_points.iterate().xyz() ) {
            double interpolated = cubic_interpolation( p, f );
            double exact        = fp( p );
            Log::info() << p << "  -->  " << interpolated << "      [exact] " << exact << std::endl;
            EXPECT( is_approximately_equal( interpolated, exact, tolerance ) );
        }
    }


    SECTION( "official version" ) {
        auto matrix_free = Config( "matrix_free", true );
        Interpolation interpolation( option::type( "tricubic" ) | matrix_free, fs, departure_points );

        Field output = Field( "output", make_datatype<double>(), make_shape( departure_points.size() ) );
        interpolation.execute( input, output );

        auto output_view = array::make_view<double, 1>( output );
        idx_t n{0};
        for ( auto p : departure_points.iterate().xyz() ) {
            double interpolated = output_view( n++ );
            double exact        = fp( p );
            Log::info() << p << "  -->  " << interpolated << "      [exact] " << exact << std::endl;
            EXPECT( is_approximately_equal( interpolated, exact, tolerance ) );
        }
    }

    SECTION( "SL-like" ) {
        auto matrix_free = Config( "matrix_free", true );

        auto dp_field = fs.createField<double>( option::variables( 3 ) );

        {
            auto iterator     = departure_points.iterate().xyz().begin();
            auto iterator_end = departure_points.iterate().xyz().end();
            auto dp           = array::make_view<double, 3>( dp_field );
            for ( idx_t n = 0; n < dp.shape( 0 ); ++n ) {
                for ( idx_t k = 0; k < dp.shape( 1 ); ++k ) {
                    PointXYZ p{0, 0, 0};
                    if ( iterator != iterator_end ) {
                        p = *iterator;
                        ++iterator;
                    }
                    dp( n, k, LON ) = p.x();
                    dp( n, k, LAT ) = p.y();
                    dp( n, k, ZZ )  = p.z();
                }
            }
        }
        Interpolation interpolation( option::type( "tricubic" ) | matrix_free, fs, dp_field );

        Field output = fs.createField<double>();
        interpolation.execute( input, output );

        auto output_view = array::make_view<double, 2>( output );


        auto iterator     = departure_points.iterate().xyz().begin();
        auto iterator_end = departure_points.iterate().xyz().end();

        for ( idx_t n = 0; n < output_view.shape( 0 ); ++n ) {
            for ( idx_t k = 0; k < output_view.shape( 1 ); ++k ) {
                PointXYZ p{0, 0, 0};
                if ( iterator != iterator_end ) {
                    p = *iterator;
                    ++iterator;

                    double interpolated = output_view( n, k );
                    double exact        = fp( p );
                    Log::info() << p << "  -->  " << interpolated << "      [exact] " << exact << std::endl;
                    EXPECT( is_approximately_equal( interpolated, exact, tolerance ) );
                }
            }
        }
    }
}

}  // namespace test
}  // namespace atlas

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