/*
 * (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. and Interpolation
 */

#pragma once

#include "StructuredInterpolation3D.h"

#include "atlas/array/ArrayView.h"
#include "atlas/field/Field.h"
#include "atlas/field/FieldSet.h"
#include "atlas/functionspace/NodeColumns.h"
#include "atlas/functionspace/PointCloud.h"
#include "atlas/functionspace/StructuredColumns.h"
#include "atlas/grid/StructuredGrid.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Log.h"
#include "atlas/runtime/Trace.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/Point.h"

namespace atlas {
namespace interpolation {
namespace method {

template <typename Kernel>
double StructuredInterpolation3D<Kernel>::convert_units_multiplier( const Field& field ) {
    std::string units = field.metadata().getString( "units", "degrees" );
    if ( units == "degrees" ) {
        return 1.;
    }
    if ( units == "radians" ) {
        return 180. / M_PI;
    }
    ATLAS_NOTIMPLEMENTED;
}

template <typename Kernel>
StructuredInterpolation3D<Kernel>::StructuredInterpolation3D( const Method::Config& config ) :
    Method( config ),
    matrix_free_{false},
    limiter_{false} {
    config.get( "matrix_free", matrix_free_ );
    config.get( "limiter", limiter_ );

    if ( not matrix_free_ ) {
        throw_NotImplemented( "Matrix-free StructuredInterpolation3D not implemented", Here() );
    }
}


template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_setup( const Grid& source, const Grid& target, const Cache& ) {
    if ( mpi::size() > 1 ) {
        ATLAS_NOTIMPLEMENTED;
    }


    ATLAS_ASSERT( StructuredGrid( source ) );
    FunctionSpace source_fs = functionspace::StructuredColumns( source, option::halo( kernel_->stencil_halo() ) );
    FunctionSpace target_fs = functionspace::PointCloud( target );

    do_setup( source_fs, target_fs );
}


template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_setup( const FunctionSpace& source, const FunctionSpace& target ) {
    ATLAS_TRACE( "StructuredInterpolation<" + Kernel::className() + ">::do_setup()" );

    source_ = source;
    target_ = target;

    if ( functionspace::PointCloud tgt = target ) {
        target_lonlat_   = tgt.lonlat();
        target_vertical_ = tgt.vertical();
        target_ghost_    = tgt.ghost();
    }
    else {
        ATLAS_NOTIMPLEMENTED;
    }

    setup( source );
}

template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_setup( const FunctionSpace& source, const Field& target ) {
    ATLAS_TRACE( "StructuredInterpolation<" + Kernel::className() + ">::do_setup(FunctionSpace source, Field target)" );

    source_ = source;

    if ( target.functionspace() ) {
        target_ = target.functionspace();
    }

    target_3d_ = target;

    setup( source );
}

template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_setup( const FunctionSpace& source, const FieldSet& target ) {
    ATLAS_TRACE( "StructuredInterpolation<" + Kernel::className() + ">::do_setup(FunctionSpace source,FieldSet target)" );

    source_ = source;

    ATLAS_ASSERT( target.size() >= 3 );
    if ( target[0].functionspace() ) {
        target_ = target[0].functionspace();
    }

    target_xyz_ = target;

    setup( source );
}

template <typename Kernel>
void StructuredInterpolation3D<Kernel>::print( std::ostream& ) const {
    ATLAS_NOTIMPLEMENTED;
}


template <typename Kernel>
void StructuredInterpolation3D<Kernel>::setup( const FunctionSpace& source ) {
    kernel_.reset( new Kernel( source, util::Config( "limiter", limiter_ ) ) );
}


template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_execute( const Field& src_field, Field& tgt_field ) const {
    FieldSet tgt( tgt_field );
    do_execute( FieldSet( src_field ), tgt );
}


template <typename Kernel>
void StructuredInterpolation3D<Kernel>::do_execute( const FieldSet& src_fields, FieldSet& tgt_fields ) const {
    if ( not matrix_free_ ) {
        Method::do_execute( src_fields, tgt_fields );
        return;
    }

    const idx_t N = src_fields.size();
    ATLAS_ASSERT( N == tgt_fields.size() );

    if ( N == 0 )
        return;

    haloExchange( src_fields );

    ATLAS_TRACE( "StructuredInterpolation<" + Kernel::className() + ">::do_execute()" );
    
    array::DataType datatype = src_fields[0].datatype();
    int rank                 = src_fields[0].rank();

    ATLAS_ASSERT( rank > 1 );

    for ( idx_t i = 0; i < N; ++i ) {
        ATLAS_ASSERT( src_fields[i].datatype() == datatype );
        ATLAS_ASSERT( src_fields[i].rank() == rank );
        ATLAS_ASSERT( tgt_fields[i].datatype() == datatype );
    }

    if ( datatype.kind() == array::DataType::KIND_REAL64 && rank == 2 ) {
        execute_impl<double, 2>( *kernel_, src_fields, tgt_fields );
    }
    if ( datatype.kind() == array::DataType::KIND_REAL32 && rank == 2 ) {
        execute_impl<float, 2>( *kernel_, src_fields, tgt_fields );
    }
    if ( datatype.kind() == array::DataType::KIND_REAL64 && rank == 3 ) {
        execute_impl<double, 3>( *kernel_, src_fields, tgt_fields );
    }
    if ( datatype.kind() == array::DataType::KIND_REAL32 && rank == 3 ) {
        execute_impl<float, 3>( *kernel_, src_fields, tgt_fields );
    }

    tgt_fields.set_dirty();
}


template <typename Kernel>
template <typename Value, int Rank>
void StructuredInterpolation3D<Kernel>::execute_impl( const Kernel& kernel, const FieldSet& src_fields,
                                                      FieldSet& tgt_fields ) const {
    const idx_t N = src_fields.size();

    auto make_src_view = [&]( const FieldSet& src_fields ) {
        std::vector<array::ArrayView<const Value, Rank> > src_view;
        src_view.reserve( N );
        for ( idx_t i = 0; i < N; ++i ) {
            src_view.emplace_back( array::make_view<const Value, Rank>( src_fields[i] ) );
        }
        return src_view;
    };

    // Assertions
    ATLAS_ASSERT( tgt_fields.size() == src_fields.size() );
    idx_t tgt_rank = -1;
    for ( auto& f : tgt_fields ) {
        if ( tgt_rank == -1 )
            tgt_rank = f.rank();
        if ( f.rank() != tgt_rank ) {
            throw_Exception( "target fields don't all have the same rank!", Here() );
        }
    }

    if ( functionspace::PointCloud( target() ) && tgt_rank == 1 ) {
        const idx_t out_npts = target_lonlat_.shape( 0 );

        const auto ghost    = array::make_view<int, 1>( target_ghost_ );
        const auto lonlat   = array::make_view<double, 2>( target_lonlat_ );
        const auto vertical = array::make_view<double, 1>( target_vertical_ );

        auto src_view = make_src_view( src_fields );

        constexpr int TargetRank = 1;
        std::vector<array::ArrayView<Value, TargetRank> > tgt_view;
        tgt_view.reserve( N );
        for ( idx_t i = 0; i < N; ++i ) {
            tgt_view.emplace_back( array::make_view<Value, TargetRank>( tgt_fields[i] ) );
        }

        const double convert_units = convert_units_multiplier( target_lonlat_ );
        atlas_omp_parallel {
            typename Kernel::Stencil stencil;
            typename Kernel::Weights weights;
            atlas_omp_for( idx_t n = 0; n < out_npts; ++n ) {
                if ( not ghost( n ) ) {
                    double x = lonlat( n, LON ) * convert_units;
                    double y = lonlat( n, LAT ) * convert_units;
                    double z = vertical( n );
                    kernel.compute_stencil( x, y, z, stencil );
                    kernel.compute_weights( x, y, z, stencil, weights );
                    for ( idx_t i = 0; i < N; ++i ) {
                        kernel.interpolate( stencil, weights, src_view[i], tgt_view[i], n );
                    }
                }
            }
        }
    }
    else if ( target_3d_ && tgt_rank == Rank ) {
        const idx_t out_npts = target_3d_.shape( 0 );
        const idx_t out_nlev = target_3d_.shape( 1 );

        const auto coords   = array::make_view<const double, 3>( target_3d_ );
        const auto src_view = make_src_view( src_fields );

        constexpr int TargetRank = Rank;
        std::vector<array::ArrayView<Value, TargetRank> > tgt_view;
        tgt_view.reserve( N );

        for ( idx_t i = 0; i < N; ++i ) {
            tgt_view.emplace_back( array::make_view<Value, TargetRank>( tgt_fields[i] ) );

            if ( Rank == 3 &&
                 ( src_fields[i].stride( Rank - 1 ) != 1 || tgt_fields[i].stride( TargetRank - 1 ) != 1 ) ) {
                throw_Exception(
                    "Something will go seriously wrong if we continue from here as "
                    "the implementation assumes stride=1 for fastest moving index (variables).",
                    Here() );
            }
        }

        const double convert_units = convert_units_multiplier( target_3d_ );

        atlas_omp_parallel {
            typename Kernel::Stencil stencil;
            typename Kernel::Weights weights;
            atlas_omp_for( idx_t n = 0; n < out_npts; ++n ) {
                for ( idx_t k = 0; k < out_nlev; ++k ) {
                    double x = coords( n, k, LON ) * convert_units;
                    double y = coords( n, k, LAT ) * convert_units;
                    double z = coords( n, k, ZZ );

                    kernel.compute_stencil( x, y, z, stencil );
                    kernel.compute_weights( x, y, z, stencil, weights );
                    for ( idx_t i = 0; i < N; ++i ) {
                        kernel.interpolate( stencil, weights, src_view[i], tgt_view[i], n, k );
                    }
                }
            }
        }
    }
    else if ( not target_xyz_.empty() && tgt_rank == Rank ) {
        const idx_t out_npts = target_xyz_[0].shape( 0 );
        const idx_t out_nlev = target_xyz_[0].shape( 1 );

        const auto xcoords  = array::make_view<double, 2>( target_xyz_[LON] );
        const auto ycoords  = array::make_view<double, 2>( target_xyz_[LAT] );
        const auto zcoords  = array::make_view<double, 2>( target_xyz_[ZZ] );
        const auto src_view = make_src_view( src_fields );

        constexpr int TargetRank = Rank;
        std::vector<array::ArrayView<Value, TargetRank> > tgt_view;
        tgt_view.reserve( N );

        for ( idx_t i = 0; i < N; ++i ) {
            tgt_view.emplace_back( array::make_view<Value, TargetRank>( tgt_fields[i] ) );

            if ( Rank == 3 &&
                 ( src_fields[i].stride( Rank - 1 ) != 1 || tgt_fields[i].stride( TargetRank - 1 ) != 1 ) ) {
                throw_Exception(
                    "Something will go seriously wrong if we continue from here as "
                    "the implementation assumes stride=1 for fastest moving index (variables).",
                    Here() );
            }
        }

        const double convert_units = convert_units_multiplier( target_xyz_[LON] );

        atlas_omp_parallel {
            typename Kernel::Stencil stencil;
            typename Kernel::Weights weights;
            atlas_omp_for( idx_t n = 0; n < out_npts; ++n ) {
                for ( idx_t k = 0; k < out_nlev; ++k ) {
                    const double x = xcoords( n, k ) * convert_units;
                    const double y = ycoords( n, k ) * convert_units;
                    const double z = zcoords( n, k );
                    kernel.compute_stencil( x, y, z, stencil );
                    kernel.compute_weights( x, y, z, stencil, weights );
                    for ( idx_t i = 0; i < N; ++i ) {
                        kernel.interpolate( stencil, weights, src_view[i], tgt_view[i], n, k );
                    }
                }
            }
        }
    }

    else {
        ATLAS_NOTIMPLEMENTED;
    }
}

}  // namespace method
}  // namespace interpolation
}  // namespace atlas
