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

#include "eckit/config/Resource.h"

#include "atlas/array/MakeView.h"
#include "atlas/field/Field.h"
#include "atlas/field/FieldSet.h"
#include "atlas/grid/Distribution.h"
#include "atlas/grid/Grid.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/meshgenerator.h"
#include "atlas/numerics/Nabla.h"
#include "atlas/numerics/fvm/Method.h"
#include "atlas/option.h"
#include "atlas/output/Gmsh.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/util/Config.h"
#include "atlas/util/Constants.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/Earth.h"

#include "tests/AtlasTestEnvironment.h"

using namespace eckit;
using namespace atlas::numerics;
using namespace atlas::meshgenerator;
using namespace atlas::grid;

namespace atlas {
namespace test {


// This test relates to JIRA issues METV-2657 , MIR-459

static std::string griduid() {
    //return "O80";
    //return "Slat80";
    return "Slat720x360";
}
static double radius() {
    return util::Earth::radius();
}
static double beta_in_degrees() {
    return 90.;
}

static bool output_gmsh() {
    return false;
}

static bool mask_polar_values() {
    return false;
}

static int metric_approach() {
    // Experimental!!!
    // approach = 0  ORIGINAL, DEFAULT
    //   metric_term cos(y) is multiplied with each wind component
    //   --> cell-interface:  avg = 0.5*( cos(y1)*u1 + cos(y2)*u2 )
    // approach = 1:
    //   metric_term cos(0.5*(y1+y2)) is used at cell interface
    //   --> cell-interface: avg = 0.5*(u1+u2)*cos(0.5*(y1+y2))
    // Results seem to indicate that approach=0 is overall better, although approach=1
    // seems to handle pole slightly better (error factor 2 to 4 times lower)
    return 0;
}

class SolidBodyRotation {
    double radius;
    double beta;
    double sin_beta;
    double cos_beta;

public:
    SolidBodyRotation(const double _radius, const double _beta): radius{_radius}, beta{_beta} {
        sin_beta = std::sin(beta);
        cos_beta = std::cos(beta);
    }
    void wind(const double x, const double y, double& u, double& v) {
        double cos_x = std::cos(x);
        double cos_y = std::cos(y);
        double sin_x = std::sin(x);
        double sin_y = std::sin(y);
        u            = cos_y * cos_beta + cos_x * sin_y * sin_beta;
        v            = -sin_x * sin_beta;
    }

    void vordiv(const double x, const double y, double& vor, double& div) {
        double cos_x = std::cos(x);
        double cos_y = std::cos(y);
        double sin_x = std::sin(x);
        double sin_y = std::sin(y);

        // Divergence = 1./(R*cos(y)) * ( d/dx( u ) + d/dy( v * cos(y) ) )
        // Vorticity  = 1./(R*cos(y)) * ( d/dx( v ) - d/dy( u * cos(y) ) )
        double ddx_u      = -sin_x * sin_y * sin_beta;
        double ddy_cosy_v = (-sin_x * sin_beta) * (-sin_y);
        double ddx_v      = -cos_x * sin_beta;
        double ddy_cosy_u =
            2 * cos_y * (-sin_y) * cos_beta + (-sin_y) * cos_x * sin_y * sin_beta + cos_y * cos_x * cos_y * sin_beta;

        double metric = 1. / (radius * cos_y);

        div = metric * (ddx_u + ddy_cosy_v);
        vor = metric * (ddx_v - ddy_cosy_u);
    }

    void wind_magnitude_squared(const double x, const double y, double& f) {
        double u, v;
        wind(x, y, u, v);
        f = u * u + v * v;
    }

    void wind_magnitude_squared_gradient(const double x, const double y, double& dfdx, double& dfdy) {
        double cos_x = std::cos(x);
        double cos_y = std::cos(y);
        double sin_x = std::sin(x);
        double sin_y = std::sin(y);

        double metric_y = 1. / radius;
        double metric_x = metric_y / cos_y;

        double u    = cos_y * cos_beta + cos_x * sin_y * sin_beta;
        double v    = -sin_x * sin_beta;
        double dudx = metric_x * (-sin_x * sin_y * sin_beta);
        double dudy = metric_y * (-sin_y * cos_beta + cos_x * cos_y * sin_beta);
        double dvdx = metric_x * (-cos_x * sin_beta);
        double dvdy = metric_y * (0.);
        dfdx        = 2 * u * dudx + 2 * v * dvdx;
        dfdy        = 2 * u * dudy + 2 * v * dvdy;
    }
};

FieldSet analytical_fields(const fvm::Method& fvm) {
    constexpr double deg2rad = M_PI / 180.;
    const double radius      = fvm.radius();

    auto lonlat_deg = array::make_view<double, 2>(fvm.mesh().nodes().lonlat());

    FieldSet fields;
    auto add_scalar_field = [&](const std::string& name) {
        return array::make_view<double, 1>(fields.add(fvm.node_columns().createField<double>(option::name(name))));
    };
    auto add_vector_field = [&](const std::string& name) {
        return array::make_view<double, 2>(fields.add(fvm.node_columns().createField<double>(
            option::name(name) | option::type("vector") | option::variables(2))));
    };
    auto f      = add_scalar_field("f");
    auto uv     = add_vector_field("uv");
    auto u      = add_scalar_field("u");
    auto v      = add_scalar_field("v");
    auto grad_f = add_vector_field("ref_grad_f");
    auto dfdx   = add_scalar_field("ref_dfdx");
    auto dfdy   = add_scalar_field("ref_dfdy");
    auto div    = add_scalar_field("ref_div");
    auto vor    = add_scalar_field("ref_vor");

    auto flow          = SolidBodyRotation{radius, beta_in_degrees() * deg2rad};
    auto is_ghost      = array::make_view<int, 1>(fvm.mesh().nodes().ghost());
    const idx_t nnodes = fvm.mesh().nodes().size();
    for (idx_t jnode = 0; jnode < nnodes; ++jnode) {
        if (is_ghost(jnode)) {
            continue;
        }
        double x = lonlat_deg(jnode, LON) * deg2rad;
        double y = lonlat_deg(jnode, LAT) * deg2rad;

        flow.wind(x, y, u(jnode), v(jnode));
        flow.vordiv(x, y, vor(jnode), div(jnode));
        flow.wind_magnitude_squared(x, y, f(jnode));
        flow.wind_magnitude_squared_gradient(x, y, dfdx(jnode), dfdy(jnode));

        uv(jnode, XX)     = u(jnode);
        uv(jnode, YY)     = v(jnode);
        grad_f(jnode, XX) = dfdx(jnode);
        grad_f(jnode, YY) = dfdy(jnode);
    }
    fields.set_dirty();
    fields.haloExchange();

    return fields;
}

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

CASE("test_analytical") {
    Grid grid(griduid(), GlobalDomain(-180.));

    Mesh mesh = MeshGenerator{"structured"}.generate(grid);
    fvm::Method fvm(mesh, option::radius(radius()));
    Nabla nabla(fvm, util::Config("metric_approach", metric_approach()));
    FieldSet fields = analytical_fields(fvm);

    Field div    = fields.add(fvm.node_columns().createField<double>(option::name("div")));
    Field vor    = fields.add(fvm.node_columns().createField<double>(option::name("vor")));
    Field grad_f = fields.add(fvm.node_columns().createField<double>(option::name("grad_f") | option::variables(2)));
    Field dfdx   = fields.add(fvm.node_columns().createField<double>(option::name("dfdx")));
    Field dfdy   = fields.add(fvm.node_columns().createField<double>(option::name("dfdy")));

    auto split = [](const Field& vector, Field& component_x, Field& component_y) {
        auto v = array::make_view<double, 2>(vector);
        auto x = array::make_view<double, 1>(component_x);
        auto y = array::make_view<double, 1>(component_y);
        for (idx_t j = 0; j < v.shape(0); ++j) {
            x(j) = v(j, XX);
            y(j) = v(j, YY);
        }
    };


    ATLAS_TRACE_SCOPE("gradient") nabla.gradient(fields["f"], grad_f);
    split(grad_f, dfdx, dfdy);
    ATLAS_TRACE_SCOPE("divergence") nabla.divergence(fields["uv"], div);
    ATLAS_TRACE_SCOPE("vorticity") nabla.curl(fields["uv"], vor);

    auto do_mask_polar_values = [&](Field& field, double mask) {
        using Topology  = atlas::mesh::Nodes::Topology;
        using Range     = atlas::array::Range;
        auto node_flags = array::make_view<int, 1>(fvm.node_columns().nodes().flags());
        auto is_polar   = [&](idx_t j) {
            return Topology::check(node_flags(j), Topology::BC | Topology::NORTH) ||
                   Topology::check(node_flags(j), Topology::BC | Topology::SOUTH);
        };
        auto apply = [&](array::LocalView<double, 2>&& view) {
            for (idx_t j = 0; j < view.shape(0); ++j) {
                if (is_polar(j)) {
                    for (idx_t v = 0; v < view.shape(1); ++v) {
                        view(j, v) = mask;
                    }
                }
            }
        };
        if (field.rank() == 1) {
            apply(array::make_view<double, 1>(field).slice(Range::all(), Range::dummy()));
        }
        else if (field.rank() == 2) {
            apply(array::make_view<double, 2>(field).slice(Range::all(), Range::all()));
        }
    };

    for (auto fieldname : std::vector<std::string>{"dfdx", "dfdy", "div", "vor"}) {
        auto err_field  = fields.add(fvm.node_columns().createField<double>(option::name("err_" + fieldname)));
        auto err2_field = fields.add(fvm.node_columns().createField<double>(option::name("err2_" + fieldname)));
        auto fld        = array::make_view<double, 1>(fields[fieldname]);
        auto ref        = array::make_view<double, 1>(fields["ref_" + fieldname]);
        auto err        = array::make_view<double, 1>(fields["err_" + fieldname]);
        auto err2       = array::make_view<double, 1>(fields["err2_" + fieldname]);
        for (idx_t j = 0; j < fld.shape(0); ++j) {
            err(j)  = fld(j) - ref(j);
            err2(j) = err(j) * err(j);
        }

        if (mask_polar_values()) {
            do_mask_polar_values(fields["err_" + fieldname], 0.);
            do_mask_polar_values(fields["err2_" + fieldname], 0.);
        }
    }

    fields.haloExchange();


    // output to gmsh
    if (output_gmsh()) {
        output::Gmsh{"mesh_2d.msh", util::Config("coordinates", "lonlat")}.write(mesh);
        output::Gmsh{"mesh_3d.msh", util::Config("coordinates", "xyz")}.write(mesh);
        output::Gmsh{"fields.msh"}.write(fields);
    }

    auto minmax_within_error = [&](const std::string& name, double error) {
        Field field = fields["err_" + name];
        error       = std::abs(error);
        double min, max;
        fvm.node_columns().minimum(field, min);
        fvm.node_columns().maximum(field, max);
        bool success = true;
        if (min < -error) {
            Log::warning() << "minumum " << min << " smaller than error " << -error << std::endl;
            success = false;
        }
        if (max > error) {
            Log::warning() << "maximum " << max << " greater than error " << error << std::endl;
            success = false;
        }
        Log::info() << name << "\t: minmax error between { " << min << " , " << max << " }" << std::endl;
        return success;
    };
    EXPECT(minmax_within_error("dfdx", 1.e-11));
    EXPECT(minmax_within_error("dfdy", 1.e-11));
    EXPECT(minmax_within_error("div", 1.e-16));
    EXPECT(minmax_within_error("vor", 1.5e-9));

    auto rms_within_error = [&](const std::string& name, double error) {
        Field field = fields["err2_" + name];
        double mean;
        idx_t N;
        fvm.node_columns().mean(field, mean, N);
        double rms   = std::sqrt(mean / double(N));
        bool success = true;
        if (rms > error) {
            Log::warning() << "rms " << rms << " greater than error " << error << std::endl;
            success = false;
        }
        Log::info() << name << "\t: rms error = " << rms << std::endl;
        return success;
    };
    EXPECT(rms_within_error("dfdx", 1.e-14));
    EXPECT(rms_within_error("dfdy", 1.e-14));
    EXPECT(rms_within_error("div", 5.e-20));
    EXPECT(rms_within_error("vor", 5.e-13));

    // error for vorticity seems too high ?
}

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

}  // namespace test
}  // namespace atlas

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