/*
 * (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 <iostream>

#include "atlas/array.h"
#include "atlas/array/ArrayView.h"
#include "atlas/array/IndexView.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/mesh/actions/BuildPeriodicBoundaries.h"
#include "atlas/parallel/mpi/Statistics.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/runtime/Exception.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/LonLatMicroDeg.h"
#include "atlas/util/PeriodicTransform.h"

using Topology = atlas::mesh::Nodes::Topology;
using atlas::util::LonLatMicroDeg;
using atlas::util::PeriodicTransform;

namespace atlas {
namespace mesh {
namespace actions {

using uid_t = gidx_t;

void build_periodic_boundaries( Mesh& mesh ) {
    ATLAS_TRACE();
    bool periodic = false;
    mesh.metadata().get( "periodic", periodic );

    auto mpi_size = mpi::size();
    auto mypart   = mpi::rank();

    if ( !periodic ) {
        mesh::Nodes& nodes = mesh.nodes();

        auto flags = array::make_view<int, 1>( nodes.flags() );
        auto ridx  = array::make_indexview<idx_t, 1>( nodes.remote_index() );
        auto part  = array::make_view<int, 1>( nodes.partition() );
        auto ghost = array::make_view<int, 1>( nodes.ghost() );

        int nb_nodes = nodes.size();

        auto xy = array::make_view<double, 2>( nodes.xy() );

        // Identify my master and slave nodes on own partition
        // master nodes are at x=0,  slave nodes are at x=2pi
        std::map<uid_t, int> master_lookup;
        std::map<uid_t, int> slave_lookup;
        std::vector<int> master_nodes;
        master_nodes.reserve( 3 * nb_nodes );
        std::vector<int> slave_nodes;
        slave_nodes.reserve( 3 * nb_nodes );

        for ( idx_t jnode = 0; jnode < nodes.size(); ++jnode ) {
            if ( Topology::check_all( flags( jnode ), Topology::BC | Topology::WEST ) ) {
                Topology::set( flags( jnode ), Topology::PERIODIC );
                if ( part( jnode ) == mypart ) {
                    LonLatMicroDeg ll( xy( jnode, XX ), xy( jnode, YY ) );
                    master_lookup[ll.unique()] = jnode;
                    master_nodes.push_back( ll.lon() );
                    master_nodes.push_back( ll.lat() );
                    master_nodes.push_back( jnode );
                }
            }
            else if ( Topology::check( flags( jnode ), Topology::BC | Topology::EAST ) ) {
                Topology::set( flags( jnode ), Topology::PERIODIC );
                Topology::set( flags( jnode ), Topology::GHOST );
                ghost( jnode ) = 1;
                LonLatMicroDeg ll( xy( jnode, XX ), xy( jnode, YY ) );
                slave_lookup[ll.unique()] = jnode;
                slave_nodes.push_back( ll.lon() );
                slave_nodes.push_back( ll.lat() );
                slave_nodes.push_back( jnode );
                ridx( jnode ) = -1;
            }
        }

        std::vector<std::vector<int>> found_master( mpi_size );
        std::vector<std::vector<int>> send_slave_idx( mpi_size );

        // Find masters on other tasks to send to me
        {
            int sendcnt = slave_nodes.size();
            std::vector<int> recvcounts( mpi_size );

            ATLAS_TRACE_MPI( ALLGATHER ) { mpi::comm().allGather( sendcnt, recvcounts.begin(), recvcounts.end() ); }

            std::vector<int> recvdispls( mpi_size );
            recvdispls[0] = 0;
            int recvcnt   = recvcounts[0];
            for ( idx_t jproc = 1; jproc < mpi_size; ++jproc ) {
                recvdispls[jproc] = recvdispls[jproc - 1] + recvcounts[jproc - 1];
                recvcnt += recvcounts[jproc];
            }
            std::vector<int> recvbuf( recvcnt );

            ATLAS_TRACE_MPI( ALLGATHER ) {
                mpi::comm().allGatherv( slave_nodes.begin(), slave_nodes.end(), recvbuf.begin(), recvcounts.data(),
                                        recvdispls.data() );
            }

            PeriodicTransform transform;
            for ( idx_t jproc = 0; jproc < mpi_size; ++jproc ) {
                found_master.reserve( master_nodes.size() );
                send_slave_idx.reserve( master_nodes.size() );
                array::LocalView<int, 2> recv_slave( recvbuf.data() + recvdispls[jproc],
                                                     array::make_shape( recvcounts[jproc] / 3, 3 ) );
                for ( idx_t jnode = 0; jnode < recv_slave.shape( 0 ); ++jnode ) {
                    LonLatMicroDeg slave( recv_slave( jnode, LON ), recv_slave( jnode, LAT ) );
                    transform( slave, -1 );
                    uid_t slave_uid = slave.unique();
                    if ( master_lookup.count( slave_uid ) ) {
                        int master_idx = master_lookup[slave_uid];
                        int slave_idx  = recv_slave( jnode, 2 );
                        found_master[jproc].push_back( master_idx );
                        send_slave_idx[jproc].push_back( slave_idx );
                    }
                }
            }
        }

        // Fill in data to communicate
        std::vector<std::vector<int>> recv_slave_idx( mpi_size );
        std::vector<std::vector<int>> send_master_part( mpi_size );
        std::vector<std::vector<int>> recv_master_part( mpi_size );
        std::vector<std::vector<int>> send_master_ridx( mpi_size );
        std::vector<std::vector<int>> recv_master_ridx( mpi_size );

        //  std::vector< std::vector<int> > send_slave_part( mpi_size );
        //  std::vector< std::vector<int> > recv_slave_part( mpi_size );
        //  std::vector< std::vector<int> > send_slave_ridx( mpi_size );
        //  std::vector< std::vector<int> > recv_slave_ridx( mpi_size );

        {
            for ( idx_t jproc = 0; jproc < mpi_size; ++jproc ) {
                idx_t nb_found_master = static_cast<idx_t>( found_master[jproc].size() );
                send_master_part[jproc].resize( nb_found_master );
                send_master_ridx[jproc].resize( nb_found_master );
                for ( idx_t jnode = 0; jnode < nb_found_master; ++jnode ) {
                    int loc_idx                    = found_master[jproc][jnode];
                    send_master_part[jproc][jnode] = part( loc_idx );
                    send_master_ridx[jproc][jnode] = loc_idx;
                }

                //      int nb_found_slaves = found_slave[jproc].size();
                //      send_slave_glb_idx[jproc].resize(nb_found_slaves);
                //      send_slave_part   [jproc].resize(nb_found_slaves);
                //      send_slave_ridx   [jproc].resize(nb_found_slaves);
                //      for( int jnode=0; jnode<nb_found_slaves; ++jnode )
                //      {
                //        int loc_idx = found_slave[jproc][jnode];
                //        send_slave_glb_idx[jproc][jnode] = glb_idx( loc_idx );
                //        send_slave_part   [jproc][jnode] = part   ( loc_idx );
                //        send_slave_ridx   [jproc][jnode] = loc_idx;
                //      }
            }
        }

        // Communicate
        ATLAS_TRACE_MPI( ALLTOALL ) {
            mpi::comm().allToAll( send_slave_idx, recv_slave_idx );
            mpi::comm().allToAll( send_master_part, recv_master_part );
            mpi::comm().allToAll( send_master_ridx, recv_master_ridx );
            //  mpi::comm().allToAll( send_slave_part,     recv_slave_part    );
            //  mpi::comm().allToAll( send_slave_loc,      recv_slave_ridx    );
        }

        // Fill in periodic
        // unused // int nb_recv_master = 0;
        for ( idx_t jproc = 0; jproc < mpi_size; ++jproc ) {
            idx_t nb_recv = static_cast<idx_t>( recv_slave_idx[jproc].size() );
            for ( idx_t jnode = 0; jnode < nb_recv; ++jnode ) {
                idx_t slave_idx   = recv_slave_idx[jproc][jnode];
                part( slave_idx ) = recv_master_part[jproc][jnode];
                ridx( slave_idx ) = recv_master_ridx[jproc][jnode];
            }
        }
    }
    mesh.metadata().set( "periodic", true );
}

// ------------------------------------------------------------------
// C wrapper interfaces to C++ routines

void atlas__build_periodic_boundaries( Mesh::Implementation* mesh ) {
    ATLAS_ASSERT( mesh != nullptr, "Cannot access uninitialised atlas_Mesh" );
    Mesh m( mesh );
    build_periodic_boundaries( m );
}
// ------------------------------------------------------------------

}  // namespace actions
}  // namespace mesh
}  // namespace atlas
