/*
 * (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 <cstdarg>
//#include <functional>

#include "eckit/utils/MD5.h"

#include "atlas/array.h"
#include "atlas/field/Field.h"
#include "atlas/field/FieldSet.h"
#include "atlas/functionspace/NodeColumns.h"
#include "atlas/grid/Grid.h"
#include "atlas/library/config.h"
#include "atlas/mesh/IsGhostNode.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/mesh/actions/BuildHalo.h"
#include "atlas/mesh/actions/BuildParallelFields.h"
#include "atlas/mesh/actions/BuildPeriodicBoundaries.h"
#include "atlas/parallel/Checksum.h"
#include "atlas/parallel/GatherScatter.h"
#include "atlas/parallel/HaloExchange.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Trace.h"
#include "atlas/util/detail/Cache.h"

#if ATLAS_HAVE_FORTRAN
#define REMOTE_IDX_BASE 1
#else
#define REMOTE_IDX_BASE 0
#endif

namespace atlas {
namespace functionspace {
namespace detail {

namespace {

template <typename T, typename Field>
array::LocalView<T, 3> make_leveled_view( Field& field ) {
    using namespace array;
    if ( field.levels() ) {
        if ( field.variables() ) {
            return make_view<T, 3>( field ).slice( Range::all(), Range::all(), Range::all() );
        }
        else {
            return make_view<T, 2>( field ).slice( Range::all(), Range::all(), Range::dummy() );
        }
    }
    else {
        if ( field.variables() ) {
            return make_view<T, 2>( field ).slice( Range::all(), Range::dummy(), Range::all() );
        }
        else {
            return make_view<T, 1>( field ).slice( Range::all(), Range::dummy(), Range::dummy() );
        }
    }
}

}  // namespace

class NodeColumnsHaloExchangeCache : public util::Cache<std::string, parallel::HaloExchange>,
                                     public mesh::detail::MeshObserver {
private:
    using Base = util::Cache<std::string, parallel::HaloExchange>;
    NodeColumnsHaloExchangeCache() : Base( "NodeColumnsHaloExchangeCache" ) {}

public:
    static NodeColumnsHaloExchangeCache& instance() {
        static NodeColumnsHaloExchangeCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const Mesh& mesh, long halo ) {
        registerMesh( *mesh.get() );
        creator_type creator = std::bind( &NodeColumnsHaloExchangeCache::create, mesh, halo );
        return Base::get_or_create( key( *mesh.get(), halo ), creator );
    }
    void onMeshDestruction( mesh::detail::MeshImpl& mesh ) override {
        for ( long jhalo = 1; jhalo <= mesh::Halo( mesh ).size(); ++jhalo ) {
            remove( key( mesh, jhalo ) );
        }
    }

private:
    static Base::key_type key( const mesh::detail::MeshImpl& mesh, long halo ) {
        std::ostringstream key;
        key << "mesh[address=" << &mesh << "],halo[size=" << halo << "]";
        return key.str();
    }

    static value_type* create( const Mesh& mesh, long halo ) {
        value_type* value = new value_type();

        std::ostringstream ss;
        ss << "nb_nodes_including_halo[" << halo << "]";
        idx_t nb_nodes( mesh.nodes().size() );
        mesh.metadata().get( ss.str(), nb_nodes );

        value->setup( array::make_view<int, 1>( mesh.nodes().partition() ).data(),
                      array::make_view<idx_t, 1>( mesh.nodes().remote_index() ).data(), REMOTE_IDX_BASE, nb_nodes );

        return value;
    }
};

class NodeColumnsGatherScatterCache : public util::Cache<std::string, parallel::GatherScatter>,
                                      public mesh::detail::MeshObserver {
private:
    using Base = util::Cache<std::string, parallel::GatherScatter>;
    NodeColumnsGatherScatterCache() : Base( "NodeColumnsGatherScatterCache" ) {}

public:
    static NodeColumnsGatherScatterCache& instance() {
        static NodeColumnsGatherScatterCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const Mesh& mesh ) {
        registerMesh( *mesh.get() );
        creator_type creator = std::bind( &NodeColumnsGatherScatterCache::create, mesh );
        return Base::get_or_create( key( *mesh.get() ), creator );
    }
    void onMeshDestruction( mesh::detail::MeshImpl& mesh ) override { remove( key( mesh ) ); }

private:
    static Base::key_type key( const mesh::detail::MeshImpl& mesh ) {
        std::ostringstream key;
        key << "mesh[address=" << &mesh << "]";
        return key.str();
    }

    static value_type* create( const Mesh& mesh ) {
        value_type* value = new value_type();

        mesh::IsGhostNode is_ghost( mesh.nodes() );
        std::vector<int> mask( mesh.nodes().size() );
        const idx_t npts = mask.size();
        atlas_omp_parallel_for( idx_t n = 0; n < npts; ++n ) {
            mask[n] = is_ghost( n ) ? 1 : 0;

            // --> This would add periodic west-bc to the gather, but means that
            // global-sums, means, etc are computed wrong
            // if( mask[j] == 1 &&
            // internals::Topology::check(flags(j),internals::Topology::BC) ) {
            //  mask[j] = 0;
            //}
        }

        value->setup( array::make_view<int, 1>( mesh.nodes().partition() ).data(),
                      array::make_view<idx_t, 1>( mesh.nodes().remote_index() ).data(), REMOTE_IDX_BASE,
                      array::make_view<gidx_t, 1>( mesh.nodes().global_index() ).data(), mask.data(),
                      mesh.nodes().size() );
        return value;
    }
};

class NodeColumnsChecksumCache : public util::Cache<std::string, parallel::Checksum>,
                                 public mesh::detail::MeshObserver {
private:
    using Base = util::Cache<std::string, parallel::Checksum>;
    NodeColumnsChecksumCache() : Base( "NodeColumnsChecksumCache" ) {}

public:
    static NodeColumnsChecksumCache& instance() {
        static NodeColumnsChecksumCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const Mesh& mesh ) {
        registerMesh( *mesh.get() );
        creator_type creator = std::bind( &NodeColumnsChecksumCache::create, mesh );
        return Base::get_or_create( key( *mesh.get() ), creator );
    }
    void onMeshDestruction( mesh::detail::MeshImpl& mesh ) override { remove( key( mesh ) ); }

private:
    static Base::key_type key( const mesh::detail::MeshImpl& mesh ) {
        std::ostringstream key;
        key << "mesh[address=" << &mesh << "]";
        return key.str();
    }

    static value_type* create( const Mesh& mesh ) {
        value_type* value = new value_type();
        util::ObjectHandle<parallel::GatherScatter> gather(
            NodeColumnsGatherScatterCache::instance().get_or_create( mesh ) );
        value->setup( gather );
        return value;
    }
};

NodeColumns::NodeColumns( Mesh mesh ) : NodeColumns( mesh, util::NoConfig() ) {}

NodeColumns::NodeColumns( Mesh mesh, const eckit::Configuration& config ) :
    mesh_( mesh ), nodes_( mesh_.nodes() ), nb_levels_( config.getInt( "levels", 0 ) ), nb_nodes_( 0 ) {
    ATLAS_TRACE();
    if ( config.has( "halo" ) ) {
        halo_ = mesh::Halo( config.getInt( "halo" ) );
    }
    else {
        halo_ = mesh::Halo( mesh );
    }
    ATLAS_ASSERT( mesh_ );
    mesh::actions::build_nodes_parallel_fields( mesh_.nodes() );
    mesh::actions::build_periodic_boundaries( mesh_ );

    if ( halo_.size() > 0 ) {
        mesh::actions::build_halo( mesh_, halo_.size() );
        std::stringstream ss;
        ss << "nb_nodes_including_halo[" << halo_.size() << "]";
        mesh_.metadata().get( ss.str(), nb_nodes_ );
    }
    if ( !nb_nodes_ ) {
        std::stringstream ss;
        ss << "nb_nodes_including_halo[" << halo_.size() << "]";
        if ( !mesh_.metadata().get( ss.str(), nb_nodes_ ) ) {
            nb_nodes_ = mesh_.nodes().size();
        }
    }
}

NodeColumns::~NodeColumns() = default;

std::string NodeColumns::distribution() const {
    return mesh().metadata().getString( "distribution" );
}

idx_t NodeColumns::nb_nodes() const {
    return nb_nodes_;
}

idx_t NodeColumns::nb_nodes_global() const {
    if ( nb_nodes_global_ >= 0 ) {
        return nb_nodes_global_;
    }
    if ( Grid grid = mesh().grid() ) {
        nb_nodes_global_ = grid.size();
    }
    else {
        nb_nodes_global_ = gather().glb_dof();
    }
    return nb_nodes_global_;
}

idx_t NodeColumns::config_nb_nodes( const eckit::Configuration& config ) const {
    idx_t size = nb_nodes();
    bool global( false );
    if ( config.get( "global", global ) ) {
        if ( global ) {
            idx_t owner( 0 );
            config.get( "owner", owner );
            idx_t _nb_nodes_global = nb_nodes_global();
            size                   = ( mpi::rank() == owner ? _nb_nodes_global : 0 );
        }
    }
    return size;
}

void NodeColumns::set_field_metadata( const eckit::Configuration& config, Field& field ) const {
    field.set_functionspace( this );

    bool global( false );
    if ( config.get( "global", global ) ) {
        if ( global ) {
            idx_t owner( 0 );
            config.get( "owner", owner );
            field.metadata().set( "owner", owner );
        }
    }
    field.metadata().set( "global", global );

    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    field.set_levels( levels );

    idx_t variables( 0 );
    config.get( "variables", variables );
    field.set_variables( variables );
}

array::DataType NodeColumns::config_datatype( const eckit::Configuration& config ) const {
    array::DataType::kind_t kind;
    if ( !config.get( "datatype", kind ) ) {
        throw_Exception( "datatype missing", Here() );
    }
    return array::DataType( kind );
}

std::string NodeColumns::config_name( const eckit::Configuration& config ) const {
    std::string name;
    config.get( "name", name );
    return name;
}

idx_t NodeColumns::config_levels( const eckit::Configuration& config ) const {
    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    return levels;
}

array::ArrayShape NodeColumns::config_shape( const eckit::Configuration& config ) const {
    array::ArrayShape shape;

    shape.push_back( config_nb_nodes( config ) );

    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    if ( levels > 0 ) {
        shape.push_back( levels );
    }

    idx_t variables( 0 );
    config.get( "variables", variables );
    if ( variables > 0 ) {
        shape.push_back( variables );
    }

    return shape;
}

Field NodeColumns::createField( const eckit::Configuration& config ) const {
    Field field = Field( config_name( config ), config_datatype( config ), config_shape( config ) );

    set_field_metadata( config, field );

    return field;
}

Field NodeColumns::createField( const Field& other, const eckit::Configuration& config ) const {
    return createField( option::datatype( other.datatype() ) | option::levels( other.levels() ) |
                        option::variables( other.variables() ) | config );
}

namespace {

template <int RANK>
void dispatch_haloExchange( Field& field, const parallel::HaloExchange& halo_exchange, bool on_device ) {
    if ( field.datatype() == array::DataType::kind<int>() ) {
        halo_exchange.template execute<int, RANK>( field.array(), on_device );
    }
    else if ( field.datatype() == array::DataType::kind<long>() ) {
        halo_exchange.template execute<long, RANK>( field.array(), on_device );
    }
    else if ( field.datatype() == array::DataType::kind<float>() ) {
        halo_exchange.template execute<float, RANK>( field.array(), on_device );
    }
    else if ( field.datatype() == array::DataType::kind<double>() ) {
        halo_exchange.template execute<double, RANK>( field.array(), on_device );
    }
    else {
        throw_Exception( "datatype not supported", Here() );
    }
    field.set_dirty( false );
}
}  // namespace

void NodeColumns::haloExchange( const FieldSet& fieldset, bool on_device ) const {
    for ( idx_t f = 0; f < fieldset.size(); ++f ) {
        Field& field = const_cast<FieldSet&>( fieldset )[f];
        switch ( field.rank() ) {
            case 1:
                dispatch_haloExchange<1>( field, halo_exchange(), on_device );
                break;
            case 2:
                dispatch_haloExchange<2>( field, halo_exchange(), on_device );
                break;
            case 3:
                dispatch_haloExchange<3>( field, halo_exchange(), on_device );
                break;
            case 4:
                dispatch_haloExchange<4>( field, halo_exchange(), on_device );
                break;
            default:
                throw_Exception( "Rank not supported", Here() );
        }
    }
}

void NodeColumns::haloExchange( const Field& field, bool on_device ) const {
    FieldSet fieldset;
    fieldset.add( field );
    haloExchange( fieldset, on_device );
}
const parallel::HaloExchange& NodeColumns::halo_exchange() const {
    if ( halo_exchange_ ) {
        return *halo_exchange_;
    }
    halo_exchange_ = NodeColumnsHaloExchangeCache::instance().get_or_create( mesh_, halo_.size() );
    return *halo_exchange_;
}

void NodeColumns::gather( const FieldSet& local_fieldset, FieldSet& global_fieldset ) const {
    ATLAS_ASSERT( local_fieldset.size() == global_fieldset.size() );

    for ( idx_t f = 0; f < local_fieldset.size(); ++f ) {
        const Field& loc      = local_fieldset[f];
        Field& glb            = global_fieldset[f];
        const idx_t nb_fields = 1;
        idx_t root( 0 );
        glb.metadata().get( "owner", root );

        if ( loc.datatype() == array::DataType::kind<int>() ) {
            parallel::Field<int const> loc_field( make_leveled_view<const int>( loc ) );
            parallel::Field<int> glb_field( make_leveled_view<int>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<long>() ) {
            parallel::Field<long const> loc_field( make_leveled_view<const long>( loc ) );
            parallel::Field<long> glb_field( make_leveled_view<long>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<float>() ) {
            parallel::Field<float const> loc_field( make_leveled_view<const float>( loc ) );
            parallel::Field<float> glb_field( make_leveled_view<float>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<double>() ) {
            parallel::Field<double const> loc_field( make_leveled_view<const double>( loc ) );
            parallel::Field<double> glb_field( make_leveled_view<double>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }
    }
}

void NodeColumns::gather( const Field& local, Field& global ) const {
    FieldSet local_fields;
    FieldSet global_fields;
    local_fields.add( local );
    global_fields.add( global );
    gather( local_fields, global_fields );
}
const parallel::GatherScatter& NodeColumns::gather() const {
    if ( gather_scatter_ ) {
        return *gather_scatter_;
    }
    gather_scatter_ = NodeColumnsGatherScatterCache::instance().get_or_create( mesh_ );
    return *gather_scatter_;
}
const parallel::GatherScatter& NodeColumns::scatter() const {
    if ( gather_scatter_ ) {
        return *gather_scatter_;
    }
    gather_scatter_ = NodeColumnsGatherScatterCache::instance().get_or_create( mesh_ );
    return *gather_scatter_;
}

void NodeColumns::scatter( const FieldSet& global_fieldset, FieldSet& local_fieldset ) const {
    ATLAS_ASSERT( local_fieldset.size() == global_fieldset.size() );

    for ( idx_t f = 0; f < local_fieldset.size(); ++f ) {
        const Field& glb      = global_fieldset[f];
        Field& loc            = local_fieldset[f];
        const idx_t nb_fields = 1;
        idx_t root( 0 );
        glb.metadata().get( "owner", root );

        if ( loc.datatype() == array::DataType::kind<int>() ) {
            parallel::Field<int const> glb_field( make_leveled_view<const int>( glb ) );
            parallel::Field<int> loc_field( make_leveled_view<int>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<long>() ) {
            parallel::Field<long const> glb_field( make_leveled_view<const long>( glb ) );
            parallel::Field<long> loc_field( make_leveled_view<long>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<float>() ) {
            parallel::Field<float const> glb_field( make_leveled_view<const float>( glb ) );
            parallel::Field<float> loc_field( make_leveled_view<float>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<double>() ) {
            parallel::Field<double const> glb_field( make_leveled_view<const double>( glb ) );
            parallel::Field<double> loc_field( make_leveled_view<double>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }

        glb.metadata().broadcast( loc.metadata(), root );
        loc.metadata().set( "global", false );
    }
}

void NodeColumns::scatter( const Field& global, Field& local ) const {
    FieldSet global_fields;
    FieldSet local_fields;
    global_fields.add( global );
    local_fields.add( local );
    scatter( global_fields, local_fields );
}

namespace {
template <typename T>
std::string checksum_3d_field( const parallel::Checksum& checksum, const Field& field ) {
    auto values = make_leveled_view<const T>( field );
    array::ArrayT<T> surface_field( values.shape( 0 ), values.shape( 2 ) );
    auto surface     = array::make_view<T, 2>( surface_field );
    const idx_t npts = values.shape( 0 );
    atlas_omp_for( idx_t n = 0; n < npts; ++n ) {
        for ( idx_t j = 0; j < surface.shape( 1 ); ++j ) {
            surface( n, j ) = 0.;
            for ( idx_t l = 0; l < values.shape( 1 ); ++l ) {
                surface( n, j ) += values( n, l, j );
            }
        }
    }
    return checksum.execute( surface.data(), surface_field.stride( 0 ) );
}
}  // namespace

std::string NodeColumns::checksum( const FieldSet& fieldset ) const {
    eckit::MD5 md5;
    for ( idx_t f = 0; f < fieldset.size(); ++f ) {
        const Field& field = fieldset[f];
        if ( field.datatype() == array::DataType::kind<int>() ) {
            md5 << checksum_3d_field<int>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<long>() ) {
            md5 << checksum_3d_field<long>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<float>() ) {
            md5 << checksum_3d_field<float>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<double>() ) {
            md5 << checksum_3d_field<double>( checksum(), field );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }
    }
    return md5;
}
std::string NodeColumns::checksum( const Field& field ) const {
    FieldSet fieldset;
    fieldset.add( field );
    return checksum( fieldset );
}

const parallel::Checksum& NodeColumns::checksum() const {
    if ( checksum_ ) {
        return *checksum_;
    }
    checksum_ = NodeColumnsChecksumCache::instance().get_or_create( mesh_ );
    return *checksum_;
}

// std::string NodesFunctionSpace::checksum( const FieldSet& fieldset ) const {
//  const parallel::Checksum& checksum = mesh_.checksum().get(checksum_name());

//  eckit::MD5 md5;
//  for( idx_t f=0; f<fieldset.size(); ++f ) {
//    const Field& field=fieldset[f];
//    if     ( field.datatype() == array::DataType::kind<int>() )
//      md5 << checksum.execute( field.data<int>(), field.stride(0) );
//    else if( field.datatype() == array::DataType::kind<long>() )
//      md5 << checksum.execute( field.data<long>(), field.stride(0) );
//    else if( field.datatype() == array::DataType::kind<float>() )
//      md5 << checksum.execute( field.data<float>(), field.stride(0) );
//    else if( field.datatype() == array::DataType::kind<double>() )
//      md5 << checksum.execute( field.data<double>(), field.stride(0) );
//    else throw_Exception("datatype not supported",Here());
//  }
//  return md5;
//}
// std::string NodesFunctionSpace::checksum( const Field& field ) const {
//  FieldSet fieldset;
//  fieldset.add(field);
//  return checksum(fieldset);
//}


}  // namespace detail

NodeColumns::NodeColumns() : FunctionSpace(), functionspace_( nullptr ) {}

NodeColumns::NodeColumns( const FunctionSpace& functionspace ) :
    FunctionSpace( functionspace ), functionspace_( dynamic_cast<const detail::NodeColumns*>( get() ) ) {}

namespace {
detail::NodeColumns* make_functionspace( Mesh mesh, const eckit::Configuration& config ) {
    return new detail::NodeColumns( mesh, config );
}
}  // namespace

NodeColumns::NodeColumns( Mesh mesh ) :
    FunctionSpace( make_functionspace( mesh, util::NoConfig() ) ),
    functionspace_( dynamic_cast<const detail::NodeColumns*>( get() ) ) {}

NodeColumns::NodeColumns( Mesh mesh, const eckit::Configuration& config ) :
    FunctionSpace( make_functionspace( mesh, config ) ),
    functionspace_( dynamic_cast<const detail::NodeColumns*>( get() ) ) {}

idx_t NodeColumns::nb_nodes() const {
    return functionspace_->nb_nodes();
}

idx_t NodeColumns::nb_nodes_global() const {  // All MPI ranks will have same output
    return functionspace_->nb_nodes_global();
}

const Mesh& NodeColumns::mesh() const {
    return functionspace_->mesh();
}

mesh::Nodes& NodeColumns::nodes() const {
    return functionspace_->nodes();
}

// -- Parallelisation aware methods

const mesh::Halo& NodeColumns::halo() const {
    return functionspace_->halo();
}

void NodeColumns::haloExchange( const FieldSet& fieldset, bool on_device ) const {
    functionspace_->haloExchange( fieldset, on_device );
}

void NodeColumns::haloExchange( const Field& field, bool on_device ) const {
    functionspace_->haloExchange( field, on_device );
}

const parallel::HaloExchange& NodeColumns::halo_exchange() const {
    return functionspace_->halo_exchange();
}

void NodeColumns::gather( const FieldSet& local, FieldSet& global ) const {
    functionspace_->gather( local, global );
}

void NodeColumns::gather( const Field& local, Field& global ) const {
    functionspace_->gather( local, global );
}

const parallel::GatherScatter& NodeColumns::gather() const {
    return functionspace_->gather();
}

void NodeColumns::scatter( const FieldSet& global, FieldSet& local ) const {
    functionspace_->scatter( global, local );
}

void NodeColumns::scatter( const Field& global, Field& local ) const {
    functionspace_->scatter( global, local );
}

const parallel::GatherScatter& NodeColumns::scatter() const {
    return functionspace_->scatter();
}

std::string NodeColumns::checksum( const FieldSet& fieldset ) const {
    return functionspace_->checksum( fieldset );
}

std::string NodeColumns::checksum( const Field& field ) const {
    return functionspace_->checksum( field );
}

const parallel::Checksum& NodeColumns::checksum() const {
    return functionspace_->checksum();
}

}  // namespace functionspace
}  // namespace atlas
