/*
 * (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 "atlas/functionspace/StructuredColumns.h"

#include <functional>
#include <iomanip>
#include <new>  // for bad_alloc exception
#include <sstream>
#include <string>
#include <unordered_map>

#include "eckit/log/Bytes.h"

#include "atlas/array/MakeView.h"
#include "atlas/grid/StructuredGrid.h"
#include "atlas/library/Library.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/parallel/GatherScatter.h"
#include "atlas/parallel/mpi/Statistics.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Log.h"
#include "atlas/runtime/Trace.h"
#include "atlas/util/vector.h"

namespace atlas {
namespace functionspace {
namespace detail {

void StructuredColumns::create_remote_index() const {
    field_remote_index_ = Field( "remote_idx", array::make_datatype<idx_t>(), array::make_shape( size_halo_ ) );
    auto remote_idx     = array::make_view<idx_t, 1>( field_remote_index_ );
    atlas_omp_parallel_for( idx_t n = 0; n < size_owned_; ++n ) { remote_idx( n ) = n; }


    ATLAS_TRACE_SCOPE( "Parallelisation ..." ) {
        auto build_partition_graph = [this]() -> std::unique_ptr<Mesh::PartitionGraph> {
            const eckit::mpi::Comm& comm = mpi::comm();
            const int mpi_size           = int( comm.size() );
            const int mpi_rank           = int( comm.rank() );

            auto part = array::make_view<int, 1>( this->partition() );

            std::vector<int> others_set( mpi_size, 0 );
            others_set[mpi_rank] = 1;
            for ( idx_t i = size_owned_; i < size_halo_; ++i ) {
                others_set[part( i )] = 1;  // present
            }
            std::vector<int> others;
            others.reserve( mpi_size );
            for ( idx_t p = 0; p < mpi_size; ++p ) {
                if ( others_set[p] ) {
                    others.emplace_back( p );
                }
            }

            eckit::mpi::Buffer<int> recv_others( mpi_size );

            ATLAS_TRACE_MPI( ALLGATHER ) { comm.allGatherv( others.begin(), others.end(), recv_others ); }

            std::vector<idx_t> counts( recv_others.counts.begin(), recv_others.counts.end() );
            std::vector<idx_t> displs( recv_others.displs.begin(), recv_others.displs.end() );
            std::vector<idx_t> values( recv_others.buffer.begin(), recv_others.buffer.end() );
            return std::unique_ptr<Mesh::PartitionGraph>(
                new Mesh::PartitionGraph( values.data(), mpi_size, displs.data(), counts.data() ) );
        };

        std::unique_ptr<Mesh::PartitionGraph> graph_ptr;
        ATLAS_TRACE_SCOPE( "Building partition graph..." ) { graph_ptr = build_partition_graph(); }
        const Mesh::PartitionGraph& graph = *graph_ptr;

        ATLAS_TRACE_SCOPE( "Setup remote_index fields..." ) {
            auto p = array::make_view<int, 1>( partition() );
            auto g = array::make_view<gidx_t, 1>( global_index() );

            const eckit::mpi::Comm& comm = mpi::comm();
            const int mpi_rank           = int( comm.rank() );

            auto neighbours           = graph.nearestNeighbours( mpi_rank );
            const idx_t nb_neighbours = static_cast<idx_t>( neighbours.size() );
            std::unordered_map<int, idx_t> part_to_neighbour;
            ATLAS_TRACE_SCOPE( "part_to_neighbour" ) {
                part_to_neighbour.reserve( nb_neighbours );
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    part_to_neighbour[neighbours[j]] = j;
                }
            }
            std::vector<idx_t> halo_per_neighbour( neighbours.size(), 0 );
            ATLAS_TRACE_SCOPE( "set halo_per_neighbour" )
            for ( idx_t i = size_owned_; i < size_halo_; ++i ) {
                halo_per_neighbour[part_to_neighbour[p( i )]]++;
            }

            std::vector<std::vector<gidx_t>> g_per_neighbour( neighbours.size() );
            ATLAS_TRACE_SCOPE( "assemble g_per_neighbour" ) {
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    g_per_neighbour[j].reserve( halo_per_neighbour[j] );
                }
                for ( idx_t j = size_owned_; j < size_halo_; ++j ) {
                    g_per_neighbour[part_to_neighbour[p( j )]].emplace_back( g( j ) );
                }
            }

            std::vector<eckit::mpi::Request> send_requests( neighbours.size() );
            std::vector<eckit::mpi::Request> recv_requests( neighbours.size() );

            std::vector<idx_t> recv_size( neighbours.size() );
            std::vector<idx_t> send_size( neighbours.size() );

            int tag = 0;
            ATLAS_TRACE_SCOPE( "send-receive g_per_neighbour size" ) {
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    send_size[j]     = static_cast<idx_t>( g_per_neighbour[j].size() );
                    send_requests[j] = comm.iSend( send_size[j], neighbours[j], tag );
                    recv_requests[j] = comm.iReceive( recv_size[j], neighbours[j], tag );
                }

                ATLAS_TRACE_MPI( WAIT ) {
                    for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                        comm.wait( send_requests[j] );
                    }

                    for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                        comm.wait( recv_requests[j] );
                    }
                }
            }

            std::vector<std::vector<gidx_t>> recv_g_per_neighbour( neighbours.size() );
            ATLAS_TRACE_SCOPE( "send-receive g_per_neighbour" )
            for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                recv_g_per_neighbour[j].resize( recv_size[j] );

                send_requests[j] =
                    comm.iSend( g_per_neighbour[j].data(), g_per_neighbour[j].size(), neighbours[j], tag );
                recv_requests[j] =
                    comm.iReceive( recv_g_per_neighbour[j].data(), recv_g_per_neighbour[j].size(), neighbours[j], tag );
            }

            std::vector<std::vector<idx_t>> send_r_per_neighbour( neighbours.size() );

            // Assemble "send_r_per_neighbour"
            // Depending if we can afford memory for a globally sized array, we can have a
            // much faster version of g_to_r map using std::vector.
            // TODO: using c++14 we can create a polymorphic lambda to avoid duplicated
            // code in the two branches of following if.

            size_t max_glb_idx = grid().size();
            atlas::vector<idx_t> g_to_r_vector;
            bool use_unordered_map_fallback = false;
            try {
                g_to_r_vector.resize( max_glb_idx + 1 );
            }
            catch ( std::bad_alloc& e ) {
                if ( comm.size() > 1 ) {
                    Log::warning() << "Could not allocate " << eckit::Bytes( ( max_glb_idx + 1 ) * sizeof( idx_t ) )
                                   << Here() << "\n"
                                   << "Using slower unordered_map fallback to map global to remote indices"
                                   << std::endl;
                    use_unordered_map_fallback = true;
                }
                else {
                    throw_Exception(
                        "Could not allocate " + std::string( eckit::Bytes( ( max_glb_idx + 1 ) * sizeof( idx_t ) ) ),
                        Here() );
                }
            }
            if ( not use_unordered_map_fallback ) {
                auto& g_to_r = g_to_r_vector;
                ATLAS_TRACE_SCOPE( "g_to_r (using vector)" ) {
                    atlas_omp_parallel_for( idx_t j = 0; j < size_owned_; ++j ) {
                        // parallel omp possible for ``` g_to_r[g(j)] = j ``` as we only loop over size_owned,
                        // where global_index is such that race-conditions cannot occur
#if ATLAS_ARRAYVIEW_BOUNDS_CHECKING
                        if ( g( j ) >= max_glb_idx + 1 ) {
                            ATLAS_DEBUG_VAR( g( j ) );
                            throw_OutOfRange( "g_to_r", g( j ), max_glb_idx + 1, Here() );
                        }
#endif
                        g_to_r[g( j )] = j;
                    }
                }
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    send_r_per_neighbour[j].reserve( recv_size[j] );

                    comm.wait( recv_requests[j] );  // wait for recv_g_per_neighbour[j]
                    for ( gidx_t gidx : recv_g_per_neighbour[j] ) {
                        send_r_per_neighbour[j].emplace_back( g_to_r[gidx] );
                    }
                }
            }
            else {
                std::unordered_map<gidx_t, idx_t> g_to_r;
                g_to_r.reserve( size_owned_ );

                ATLAS_TRACE_SCOPE( "g_to_r (using unordered set)" ) {
                    for ( idx_t j = 0; j < size_owned_; ++j ) {
                        g_to_r[g( j )] = j;
                    }
                }
                ATLAS_TRACE_SCOPE( "assemble r_per_neighbour" )
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    send_r_per_neighbour[j].reserve( recv_size[j] );

                    comm.wait( recv_requests[j] );  // wait for recv_g_per_neighbour[j]
                    for ( gidx_t gidx : recv_g_per_neighbour[j] ) {
                        send_r_per_neighbour[j].emplace_back( g_to_r[gidx] );
                    }
                }
            }

            std::vector<std::vector<idx_t>> r_per_neighbour( neighbours.size() );

            ATLAS_TRACE_SCOPE( "send-receive r_per_neighbour" ) {
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    r_per_neighbour[j].resize( halo_per_neighbour[j] );
                }

                ATLAS_TRACE_MPI( ALLTOALL ) {
                    for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                        comm.wait( send_requests[j] );
                        send_requests[j] = comm.iSend( send_r_per_neighbour[j].data(), send_r_per_neighbour[j].size(),
                                                       neighbours[j], tag );
                        recv_requests[j] =
                            comm.iReceive( r_per_neighbour[j].data(), r_per_neighbour[j].size(), neighbours[j], tag );
                    }
                }
                ATLAS_TRACE_MPI( WAIT ) {
                    for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                        comm.wait( recv_requests[j] );
                    }
                }
            }


            std::vector<idx_t> counters( neighbours.size(), 0 );
            ATLAS_TRACE_SCOPE( "set remote_idx" )
            for ( idx_t j = size_owned_; j < size_halo_; ++j ) {
                idx_t neighbour = part_to_neighbour[p( j )];
                remote_idx( j ) = r_per_neighbour[neighbour][counters[neighbour]++];
            }

            ATLAS_TRACE_MPI( WAIT ) {
                for ( idx_t j = 0; j < nb_neighbours; ++j ) {
                    comm.wait( send_requests[j] );
                }
            }
        }
    }
}

}  // namespace detail

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

}  // namespace functionspace
}  // namespace atlas
