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

#include "eckit/types/FloatCompare.h"

#include "atlas/array/ArrayView.h"
#include "atlas/array/MakeView.h"
#include "atlas/field/Field.h"
#include "atlas/grid.h"
#include "atlas/grid/detail/partitioner/EqualRegionsPartitioner.h"
#include "atlas/grid/detail/spacing/gaussian/Latitudes.h"
#include "atlas/library/config.h"
#include "atlas/mesh/Elements.h"
#include "atlas/mesh/HybridElements.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/mesh/actions/BuildParallelFields.h"
#include "atlas/meshgenerator.h"
#include "atlas/output/Gmsh.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/runtime/Log.h"
#include "atlas/util/Config.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/Metadata.h"

#include "tests/AtlasTestEnvironment.h"

namespace atlas {
namespace grid {
namespace spacing {
namespace gaussian {
void compute_gaussian_quadrature_npole_equator(const size_t N, double lats[], double weights[]);
}
}  // namespace spacing
}  // namespace grid
}  // namespace atlas

#define DISABLE if (0)
#define ENABLE if (1)

namespace atlas {
namespace test {

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

static ReducedGaussianGrid debug_grid() {
    return {6, 10, 18, 22, 22, 22, 22, 18, 10, 6};
}

static StructuredGrid minimal_grid(int N, long lon[]) {
    std::vector<long> nx(2 * N);
    for (long j = 0; j < N; ++j) {
        nx[j]                 = lon[j];
        nx[nx.size() - 1 - j] = nx[j];
    }
    return ReducedGaussianGrid(nx);
}

double compute_lonlat_area(Mesh& mesh) {
    mesh::Nodes& nodes                 = mesh.nodes();
    mesh::Elements& quads              = mesh.cells().elements(0);
    mesh::Elements& triags             = mesh.cells().elements(1);
    array::ArrayView<double, 2> lonlat = array::make_view<double, 2>(nodes.lonlat());

    const mesh::BlockConnectivity& quad_nodes  = quads.node_connectivity();
    const mesh::BlockConnectivity& triag_nodes = triags.node_connectivity();

    double area = 0;
    for (idx_t e = 0; e < quads.size(); ++e) {
        idx_t n0  = quad_nodes(e, 0);
        idx_t n1  = quad_nodes(e, 1);
        idx_t n2  = quad_nodes(e, 2);
        idx_t n3  = quad_nodes(e, 3);
        double x0 = lonlat(n0, LON), x1 = lonlat(n1, LON), x2 = lonlat(n2, LON), x3 = lonlat(n3, LON);
        double y0 = lonlat(n0, LAT), y1 = lonlat(n1, LAT), y2 = lonlat(n2, LAT), y3 = lonlat(n3, LAT);
        area += std::abs(x0 * (y1 - y2) + x1 * (y2 - y0) + x2 * (y0 - y1)) * 0.5;
        area += std::abs(x2 * (y3 - y0) + x3 * (y0 - y2) + x0 * (y2 - y3)) * 0.5;
    }
    for (idx_t e = 0; e < triags.size(); ++e) {
        idx_t n0  = triag_nodes(e, 0);
        idx_t n1  = triag_nodes(e, 1);
        idx_t n2  = triag_nodes(e, 2);
        double x0 = lonlat(n0, LON), x1 = lonlat(n1, LON), x2 = lonlat(n2, LON);
        double y0 = lonlat(n0, LAT), y1 = lonlat(n1, LAT), y2 = lonlat(n2, LAT);
        area += std::abs(x0 * (y1 - y2) + x1 * (y2 - y0) + x2 * (y0 - y1)) * 0.5;
    }
    return area;
}

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

CASE("test_eq_caps") {
    std::vector<int> n_regions;
    std::vector<double> s_cap;

    grid::detail::partitioner::eq_caps(6, n_regions, s_cap);
    EXPECT(n_regions.size() == 3);
    EXPECT(n_regions[0] == 1);
    EXPECT(n_regions[1] == 4);
    EXPECT(n_regions[2] == 1);

    grid::detail::partitioner::eq_caps(10, n_regions, s_cap);
    EXPECT(n_regions.size() == 4);
    EXPECT(n_regions[0] == 1);
    EXPECT(n_regions[1] == 4);
    EXPECT(n_regions[2] == 4);
    EXPECT(n_regions[3] == 1);
}

CASE("test_partitioner") {
    Grid g("S4x2");

    // 12 partitions
    {
        grid::detail::partitioner::EqualRegionsPartitioner partitioner(12);
        EXPECT(partitioner.nb_bands() == 4);
        EXPECT(partitioner.nb_regions(0) == 1);
        EXPECT(partitioner.nb_regions(1) == 5);
        EXPECT(partitioner.nb_regions(2) == 5);
        EXPECT(partitioner.nb_regions(3) == 1);
    }

    // 24 partitions
    {
        grid::detail::partitioner::EqualRegionsPartitioner partitioner(24);
        EXPECT(partitioner.nb_bands() == 5);
        EXPECT(partitioner.nb_regions(0) == 1);
        EXPECT(partitioner.nb_regions(1) == 6);
        EXPECT(partitioner.nb_regions(2) == 10);
        EXPECT(partitioner.nb_regions(3) == 6);
        EXPECT(partitioner.nb_regions(4) == 1);
    }

    // 48 partitions
    {
        grid::detail::partitioner::EqualRegionsPartitioner partitioner(48);
        EXPECT(partitioner.nb_bands() == 7);
        EXPECT(partitioner.nb_regions(0) == 1);
        EXPECT(partitioner.nb_regions(1) == 6);
        EXPECT(partitioner.nb_regions(2) == 11);
        EXPECT(partitioner.nb_regions(3) == 12);
        EXPECT(partitioner.nb_regions(4) == 11);
        EXPECT(partitioner.nb_regions(5) == 6);
        EXPECT(partitioner.nb_regions(6) == 1);
    }

    // 96 partitions
    {
        grid::detail::partitioner::EqualRegionsPartitioner partitioner(96);
        EXPECT(partitioner.nb_bands() == 10);
        EXPECT(partitioner.nb_regions(0) == 1);
        EXPECT(partitioner.nb_regions(1) == 6);
        EXPECT(partitioner.nb_regions(2) == 11);
        EXPECT(partitioner.nb_regions(3) == 14);
        EXPECT(partitioner.nb_regions(4) == 16);
        EXPECT(partitioner.nb_regions(5) == 16);
        EXPECT(partitioner.nb_regions(6) == 14);
        EXPECT(partitioner.nb_regions(7) == 11);
        EXPECT(partitioner.nb_regions(8) == 6);
        EXPECT(partitioner.nb_regions(9) == 1);
    }
}

CASE("test_gaussian_latitudes") {
    std::vector<double> factory_latitudes;
    std::vector<double> computed_latitudes;
    std::vector<double> computed_weights;

    size_t size_test_N = 19;

    size_t test_N[] = {16,  24,  32,  48,  64,  80,   96,   128,  160,  200,  256, 320,
                       400, 512, 576, 640, 800, 1024, 1280, 1600, 2000, 4000, 8000};

    for (size_t i = 0; i < size_test_N; ++i) {
        size_t N = test_N[i];
        Log::info() << "Testing gaussian latitude " << N << std::endl;
        factory_latitudes.resize(N);
        computed_latitudes.resize(N);
        computed_weights.resize(N);
        // grid::gaussian::latitudes::gaussian_latitudes_npole_equator (N,
        // factory_latitudes.data());
        // grid::gaussian::latitudes::compute_gaussian_quadrature_npole_equator(N,
        // computed_latitudes.data(), computed_weights.data());
        grid::spacing::gaussian::gaussian_latitudes_npole_equator(N, factory_latitudes.data());
        grid::spacing::gaussian::compute_gaussian_quadrature_npole_equator(N, computed_latitudes.data(),
                                                                           computed_weights.data());
        double wsum = 0;
        for (size_t i = 0; i < N; ++i) {
            EXPECT(eckit::types::is_approximately_equal(computed_latitudes[i], factory_latitudes[i], 0.0000001));
            wsum += computed_weights[i];
        }
        EXPECT(eckit::types::is_approximately_equal(wsum * 2., 1., 0.0000001));
    }
}

CASE("test_rgg_meshgen_one_part") {
    Mesh m;
    util::Config default_opts;
    default_opts.set("nb_parts", 1);
    default_opts.set("part", 0);
    //  generate.options.set("nb_parts",1);
    //  generate.options.set("part",    0);
    DISABLE {  // This is all valid for meshes generated with MINIMAL NB TRIAGS
        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", true)("include_pole", false));
            m = generate(atlas::test::debug_grid());
            EXPECT(m.nodes().size() == 156);
            EXPECT(m.cells().elements(0).size() == 134);
            EXPECT(m.cells().elements(1).size() == 32);
            //    EXPECT( m.nodes().metadata().get<size_t>("nb_owned") ==    156 );
            //    EXPECT( m.function_space("quads"
            //    ).metadata().get<size_t>("nb_owned") ==    134 );
            //    EXPECT(
            //    m.function_space("triags").metadata().get<size_t>("nb_owned") ==
            //    32 );
        }

        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", false)("include_pole", false));
            m = generate(atlas::test::debug_grid());
            EXPECT(m.nodes().size() == 166);
            EXPECT(m.cells().elements(0).size() == 134);
            EXPECT(m.cells().elements(1).size() == 32);
            //    EXPECT( m.nodes().metadata().get<size_t>("nb_owned") ==    166 );
            //    EXPECT( m.function_space("quads"
            //    ).metadata().get<size_t>("nb_owned") ==    134 );
            //    EXPECT(
            //    m.function_space("triags").metadata().get<size_t>("nb_owned") ==
            //    32 );
        }

        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", true)("include_pole", true));
            m = generate(atlas::test::debug_grid());
            EXPECT(m.nodes().size() == 158);
            EXPECT(m.cells().elements(0).size() == 134);
            EXPECT(m.cells().elements(1).size() == 44);
            //    EXPECT( m.nodes().metadata().get<size_t>("nb_owned") ==    158 );
            //    EXPECT( m.function_space("quads"
            //    ).metadata().get<size_t>("nb_owned") ==    134 );
            //    EXPECT(
            //    m.function_space("triags").metadata().get<size_t>("nb_owned") ==
            //    44 );
        }

        Mesh mesh;

        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", false)("include_pole", false));
            int nlat   = 2;
            long lon[] = {4, 6};
            mesh       = generate(test::minimal_grid(nlat, lon));
            EXPECT(mesh.nodes().size() == 24);
            EXPECT(mesh.cells().elements(0).size() == 14);
            EXPECT(mesh.cells().elements(1).size() == 4);

            double max_lat = test::minimal_grid(nlat, lon).y().front();
            EXPECT(
                eckit::types::is_approximately_equal(test::compute_lonlat_area(mesh), 2. * M_PI * 2. * max_lat, 1e-8));
            output::Gmsh("minimal2.msh").write(mesh);
        }
        // 3 latitudes
        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", false)("include_pole", false));
            int nlat   = 3;
            long lon[] = {4, 6, 8};
            mesh       = generate(test::minimal_grid(nlat, lon));
            EXPECT(mesh.nodes().size() == 42);
            EXPECT(mesh.cells().elements(0).size() == 28);
            EXPECT(mesh.cells().elements(1).size() == 8);
            output::Gmsh("minimal3.msh").write(mesh);
        }
        // 4 latitudes
        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", false)("include_pole", false));
            int nlat   = 4;
            long lon[] = {4, 6, 8, 10};
            mesh       = generate(test::minimal_grid(nlat, lon));
            EXPECT(mesh.nodes().size() == 64);
            EXPECT(mesh.cells().elements(0).size() == 46);
            EXPECT(mesh.cells().elements(1).size() == 12);
            output::Gmsh("minimal4.msh").write(mesh);
        }
        // 5 latitudes WIP
        ENABLE {
            StructuredMeshGenerator generate(default_opts("3d", false)("include_pole", false));
            int nlat   = 5;
            long lon[] = {6, 10, 18, 22, 22};
            mesh       = generate(test::minimal_grid(nlat, lon));
            EXPECT(mesh.nodes().size() == 166);
            EXPECT(mesh.cells().elements(0).size() == 134);
            EXPECT(mesh.cells().elements(1).size() == 32);
            output::Gmsh("minimal5.msh").write(mesh);
        }
    }
}

CASE("test_rgg_meshgen_many_parts") {
    EXPECT(grid::detail::partitioner::PartitionerFactory::has("equal_regions"));
    int nb_parts = 20;
    //  Alternative grid for debugging
    //  int nlat=10;
    //  long lon[] = { 10, 10, 10, 10, 10, 10, 10, 10, 10, 10 };
    //  test::MinimalGrid grid(nlat,lon);
    StructuredGrid grid = Grid("N32");
    // RegularGrid grid(128,64);

    /*
std::cout << grid.spec() << std::endl;
for (int jlat=0;jlat<2*nlat; jlat++) {
      std::cout << grid.lon(jlat,0) << ", ... , " << "," << grid.lon(jlat,9) <<
grid.lon(jlat,10) << std::endl;
}
ASSERT(0);
*/

    double max_lat    = grid.y().front();
    double check_area = 360. * 2. * max_lat;
    double area       = 0;
    int nodes[]  = {313, 332, 336, 338, 334, 337, 348, 359, 360, 361, 360, 360, 359, 370, 321, 334, 338, 335, 348, 315};
    int quads[]  = {243, 290, 293, 294, 291, 293, 310, 320, 321, 321, 320, 321, 320, 331, 278, 291, 294, 293, 305, 245};
    int triags[] = {42, 13, 12, 13, 12, 14, 0, 1, 0, 1, 1, 0, 1, 0, 14, 12, 13, 11, 14, 42};
    int nb_owned = 0;

    std::vector<int> all_owned(grid.size(), -1);

    for (int p = 0; p < nb_parts; ++p) {
        ATLAS_DEBUG_VAR(p);

        StructuredMeshGenerator generate(util::Config("partitioner", "equal_regions")("nb_parts", nb_parts)("part", p)(
            "include_pole", false)("3d", false));
        ATLAS_DEBUG_HERE();

        Mesh m = generate(grid);
        ATLAS_DEBUG_HERE();
        m.metadata().set("part", p);
        Log::info() << "generated grid " << p << std::endl;
        array::ArrayView<int, 1> part    = array::make_view<int, 1>(m.nodes().partition());
        array::ArrayView<gidx_t, 1> gidx = array::make_view<gidx_t, 1>(m.nodes().global_index());

        area += test::compute_lonlat_area(m);
        ATLAS_DEBUG_HERE();

        DISABLE {  // This is all valid for meshes generated with MINIMAL NB TRIAGS
            if (nb_parts == 20) {
                EXPECT(m.nodes().size() == nodes[p]);
                EXPECT(m.cells().elements(0).size() == quads[p]);
                EXPECT(m.cells().elements(1).size() == triags[p]);
            }
        }
        ATLAS_DEBUG_HERE();

        output::Gmsh("T63.msh").write(m);

        mesh::Nodes& nodes = m.nodes();
        idx_t nb_nodes     = nodes.size();

        // Test if all nodes are connected
        {
            std::vector<int> node_elem_connections(nb_nodes, 0);

            const mesh::HybridElements::Connectivity& cell_node_connectivity = m.cells().node_connectivity();
            for (idx_t jelem = 0; jelem < static_cast<idx_t>(m.cells().size()); ++jelem) {
                for (idx_t jnode = 0; jnode < cell_node_connectivity.cols(jelem); ++jnode) {
                    node_elem_connections[cell_node_connectivity(jelem, jnode)]++;
                }
            }
            for (idx_t jnode = 0; jnode < nb_nodes; ++jnode) {
                if (node_elem_connections[jnode] == 0) {
                    std::stringstream ss;
                    ss << "part " << p << ": node_gid " << gidx(jnode) << " is not connected to any element.";
                    DISABLE { Log::error() << ss.str() << std::endl; }
                }
            }
        }

        // Test if all nodes are owned
        for (idx_t n = 0; n < nb_nodes; ++n) {
            if (gidx(n) <= grid.size()) {
                if (part(n) == p) {
                    ++nb_owned;
                    if (all_owned[gidx(n) - 1] != -1) {
                        std::cout << "node " << gidx(n) << " already visited by " << all_owned[gidx(n) - 1]
                                  << std::endl;
                    }
                    EXPECT(all_owned[gidx(n) - 1] == -1);
                    all_owned[gidx(n) - 1] = part(n);
                }
            }
        }
    }

    for (size_t gid = 1; gid <= all_owned.size(); ++gid) {
        if (all_owned[gid - 1] == -1) {
            Log::error() << "node " << gid << " is not owned by anyone" << std::endl;
        }
    }
    EXPECT_EQ(nb_owned, grid.size());

    EXPECT_APPROX_EQ(area, check_area, 1e-9);
}

CASE("test_meshgen_ghost_at_end") {
    ATLAS_DEBUG_HERE();

    Grid grid("O8");

    atlas::util::Config cfg;
    cfg.set("part", 1);
    cfg.set("nb_parts", 8);
    StructuredMeshGenerator meshgenerator(cfg);
    Mesh mesh        = meshgenerator.generate(grid);
    const auto part  = array::make_view<int, 1>(mesh.nodes().partition());
    const auto ghost = array::make_view<int, 1>(mesh.nodes().ghost());
    const auto flags = array::make_view<int, 1>(mesh.nodes().flags());

    Log::info() << "partition = [ ";
    for (idx_t jnode = 0; jnode < part.size(); ++jnode) {
        Log::info() << part(jnode) << " ";
    }
    Log::info() << "]" << std::endl;

    Log::info() << "ghost     = [ ";
    for (idx_t jnode = 0; jnode < part.size(); ++jnode) {
        Log::info() << ghost(jnode) << " ";
    }
    Log::info() << "]" << std::endl;

    Log::info() << "flags     = [ ";
    for (idx_t jnode = 0; jnode < part.size(); ++jnode) {
        Log::info() << mesh::Nodes::Topology::check(flags(jnode), mesh::Nodes::Topology::GHOST) << " ";
        EXPECT(mesh::Nodes::Topology::check(flags(jnode), mesh::Nodes::Topology::GHOST) == ghost(jnode));
    }
    Log::info() << "]" << std::endl;
}

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

}  // namespace test
}  // namespace atlas

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