/*
 * (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 "eckit/config/Parametrisation.h"

#include "atlas/array/ArrayView.h"
#include "atlas/array/MakeView.h"
#include "atlas/field/Field.h"
#include "atlas/mesh/HybridElements.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/mesh/Nodes.h"
#include "atlas/numerics/fvm/Method.h"
#include "atlas/numerics/fvm/Nabla.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Log.h"
#include "atlas/util/CoordinateEnums.h"

// =======================================================

using Topology = atlas::mesh::Nodes::Topology;
using Range    = atlas::array::Range;

namespace atlas {
namespace numerics {
namespace fvm {

namespace {
static NablaBuilder<Nabla> __fvm_nabla("fvm");
}

Nabla::Nabla(const numerics::Method& method, const eckit::Parametrisation& p): atlas::numerics::NablaImpl(method, p) {
    fvm_ = dynamic_cast<const fvm::Method*>(&method);
    if (!fvm_) {
        throw_Exception("atlas::numerics::fvm::Nabla needs a atlas::numerics::fvm::Method", Here());
    }
    Log::debug() << "Nabla constructed for method " << fvm_->name() << " with "
                 << fvm_->node_columns().nb_nodes_global() << " nodes total" << std::endl;
    fvm_->attach();

    p.get("metric_approach", metric_approach_);
    // Experimental, only affects divergence and curl !!!
    // metric_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 )
    // metric_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)

    setup();
}

Nabla::~Nabla() = default;

void Nabla::setup() {
    const mesh::Edges& edges = fvm_->mesh().edges();

    const idx_t nedges = fvm_->edge_columns().nb_edges();

    const auto edge_flags = array::make_view<int, 1>(edges.flags());
    auto is_pole_edge     = [&](idx_t e) { return Topology::check(edge_flags(e), Topology::POLE); };

    // Filter pole_edges out of all edges
    std::vector<idx_t> tmp(nedges);
    idx_t c(0);
    for (idx_t jedge = 0; jedge < nedges; ++jedge) {
        if (is_pole_edge(jedge)) {
            tmp[c++] = jedge;
        }
    }
    pole_edges_.clear();
    pole_edges_.reserve(c);
    for (idx_t jedge = 0; jedge < c; ++jedge) {
        pole_edges_.push_back(tmp[jedge]);
    }
}

void Nabla::gradient(const Field& field, Field& grad_field) const {
    if (field.variables() > 1) {
        return gradient_of_vector(field, grad_field);
    }
    else {
        return gradient_of_scalar(field, grad_field);
    }
}

void Nabla::gradient_of_scalar(const Field& scalar_field, Field& grad_field) const {
    Log::debug() << "Compute gradient of scalar field " << scalar_field.name() << " with fvm method" << std::endl;
    const double radius  = fvm_->radius();
    const double deg2rad = M_PI / 180.;

    const mesh::Edges& edges = fvm_->mesh().edges();
    const mesh::Nodes& nodes = fvm_->mesh().nodes();

    const idx_t nnodes = fvm_->node_columns().nb_nodes();
    const idx_t nedges = fvm_->edge_columns().nb_edges();

    const auto scalar = scalar_field.levels()
                            ? array::make_view<double, 2>(scalar_field).slice(Range::all(), Range::all())
                            : array::make_view<double, 1>(scalar_field).slice(Range::all(), Range::dummy());
    auto grad         = grad_field.levels()
                            ? array::make_view<double, 3>(grad_field).slice(Range::all(), Range::all(), Range::all())
                            : array::make_view<double, 2>(grad_field).slice(Range::all(), Range::dummy(), Range::all());

    const idx_t nlev = scalar.shape(1);
    if (grad.shape(1) != nlev) {
        throw_AssertionFailed("gradient field should have same number of levels", Here());
    }

    const auto lonlat_deg     = array::make_view<double, 2>(nodes.lonlat());
    const auto dual_volumes   = array::make_view<double, 1>(nodes.field("dual_volumes"));
    const auto dual_normals   = array::make_view<double, 2>(edges.field("dual_normals"));
    const auto node2edge_sign = array::make_view<double, 2>(nodes.field("node2edge_sign"));

    const mesh::Connectivity& node2edge           = nodes.edge_connectivity();
    const mesh::MultiBlockConnectivity& edge2node = edges.node_connectivity();

    array::ArrayT<double> avgS_arr(nedges, nlev, 2ul);
    auto avgS = array::make_view<double, 3>(avgS_arr);

    const double scale = deg2rad * deg2rad * radius;

    atlas_omp_parallel {
        atlas_omp_for(idx_t jedge = 0; jedge < nedges; ++jedge) {
            idx_t ip1 = edge2node(jedge, 0);
            idx_t ip2 = edge2node(jedge, 1);

            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                double avg             = (scalar(ip1, jlev) + scalar(ip2, jlev)) * 0.5;
                avgS(jedge, jlev, LON) = dual_normals(jedge, LON) * deg2rad * avg;
                avgS(jedge, jlev, LAT) = dual_normals(jedge, LAT) * deg2rad * avg;
            }
        }

        atlas_omp_for(idx_t jnode = 0; jnode < nnodes; ++jnode) {
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                grad(jnode, jlev, LON) = 0.;
                grad(jnode, jlev, LAT) = 0.;
            }
            for (idx_t jedge = 0; jedge < node2edge.cols(jnode); ++jedge) {
                const idx_t iedge = node2edge(jnode, jedge);
                if (iedge < nedges) {
                    const double add = node2edge_sign(jnode, jedge);
                    for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                        grad(jnode, jlev, LON) += add * avgS(iedge, jlev, LON);
                        grad(jnode, jlev, LAT) += add * avgS(iedge, jlev, LAT);
                    }
                }
            }

            const double y        = lonlat_deg(jnode, LAT) * deg2rad;
            const double metric_y = 1. / (dual_volumes(jnode) * scale);
            const double metric_x = metric_y / std::cos(y);
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                grad(jnode, jlev, LON) *= metric_x;
                grad(jnode, jlev, LAT) *= metric_y;
            }
        }
    }
}

// ================================================================================

void Nabla::gradient_of_vector(const Field& vector_field, Field& grad_field) const {
    Log::debug() << "Compute gradient of vector field " << vector_field.name() << " with fvm method" << std::endl;
    const double radius  = fvm_->radius();
    const double deg2rad = M_PI / 180.;

    const mesh::Edges& edges = fvm_->mesh().edges();
    const mesh::Nodes& nodes = fvm_->mesh().nodes();

    const idx_t nnodes = fvm_->node_columns().nb_nodes();
    const idx_t nedges = fvm_->edge_columns().nb_edges();

    const auto vector =
        vector_field.levels()
            ? array::make_view<double, 3>(vector_field).slice(Range::all(), Range::all(), Range::all())
            : array::make_view<double, 2>(vector_field).slice(Range::all(), Range::dummy(), Range::all());
    auto grad = grad_field.levels()
                    ? array::make_view<double, 3>(grad_field).slice(Range::all(), Range::all(), Range::all())
                    : array::make_view<double, 2>(grad_field).slice(Range::all(), Range::dummy(), Range::all());

    const idx_t nlev = vector.shape(1);
    if (grad.shape(1) != nlev) {
        throw_AssertionFailed("gradient field should have same number of levels", Here());
    }

    const auto lonlat_deg     = array::make_view<double, 2>(nodes.lonlat());
    const auto dual_volumes   = array::make_view<double, 1>(nodes.field("dual_volumes"));
    const auto dual_normals   = array::make_view<double, 2>(edges.field("dual_normals"));
    const auto node2edge_sign = array::make_view<double, 2>(nodes.field("node2edge_sign"));
    const auto edge_flags     = array::make_view<int, 1>(edges.flags());
    auto is_pole_edge         = [&](idx_t e) { return Topology::check(edge_flags(e), Topology::POLE); };

    const mesh::Connectivity& node2edge           = nodes.edge_connectivity();
    const mesh::MultiBlockConnectivity& edge2node = edges.node_connectivity();

    array::ArrayT<double> avgS_arr(nedges, nlev, 4ul);
    array::ArrayView<double, 3> avgS = array::make_view<double, 3>(avgS_arr);

    const double scale = deg2rad * deg2rad * radius;

    enum
    {
        LONdLON = 0,
        LONdLAT = 1,
        LATdLON = 2,
        LATdLAT = 3
    };

    atlas_omp_parallel {
        atlas_omp_for(idx_t jedge = 0; jedge < nedges; ++jedge) {
            idx_t ip1  = edge2node(jedge, 0);
            idx_t ip2  = edge2node(jedge, 1);
            double pbc = 1. - 2. * is_pole_edge(jedge);

            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                double avg[2]              = {(vector(ip1, jlev, LON) + pbc * vector(ip2, jlev, LON)) * 0.5,
                                 (vector(ip1, jlev, LAT) + pbc * vector(ip2, jlev, LAT)) * 0.5};
                avgS(jedge, jlev, LONdLON) = dual_normals(jedge, LON) * deg2rad * avg[LON];
                // above = 0 at pole because of dual_normals
                avgS(jedge, jlev, LONdLAT) = dual_normals(jedge, LAT) * deg2rad * avg[LON];
                avgS(jedge, jlev, LATdLON) = dual_normals(jedge, LON) * deg2rad * avg[LAT];
                // above = 0 at pole because of dual_normals
                avgS(jedge, jlev, LATdLAT) = dual_normals(jedge, LAT) * deg2rad * avg[LAT];
            }
        }

        atlas_omp_for(idx_t jnode = 0; jnode < nnodes; ++jnode) {
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                grad(jnode, jlev, LONdLON) = 0.;
                grad(jnode, jlev, LONdLAT) = 0.;
                grad(jnode, jlev, LATdLON) = 0.;
                grad(jnode, jlev, LATdLAT) = 0.;
            }
            for (idx_t jedge = 0; jedge < node2edge.cols(jnode); ++jedge) {
                const idx_t iedge = node2edge(jnode, jedge);
                if (iedge < nedges) {
                    double add = node2edge_sign(jnode, jedge);
                    for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                        grad(jnode, jlev, LONdLON) += add * avgS(iedge, jlev, LONdLON);
                        grad(jnode, jlev, LONdLAT) += add * avgS(iedge, jlev, LONdLAT);
                        grad(jnode, jlev, LATdLON) += add * avgS(iedge, jlev, LATdLON);
                        grad(jnode, jlev, LATdLAT) += add * avgS(iedge, jlev, LATdLAT);
                    }
                }
            }
            const double y        = lonlat_deg(jnode, LAT) * deg2rad;
            const double metric_y = 1. / (dual_volumes(jnode) * scale);
            const double metric_x = metric_y / std::cos(y);
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                grad(jnode, jlev, LONdLON) *= metric_x;
                grad(jnode, jlev, LATdLON) *= metric_x;
                grad(jnode, jlev, LONdLAT) *= metric_y;
                grad(jnode, jlev, LATdLAT) *= metric_y;
            }
        }
    }
    // Fix wrong node2edge_sign for vector quantities
    for (size_t jedge = 0; jedge < pole_edges_.size(); ++jedge) {
        const idx_t iedge     = pole_edges_[jedge];
        const idx_t jnode     = edge2node(iedge, 1);
        const double metric_y = 1. / (dual_volumes(jnode) * scale);
        for (idx_t jlev = 0; jlev < nlev; ++jlev) {
            grad(jnode, jlev, LONdLAT) -= 2. * avgS(iedge, jlev, LONdLAT) * metric_y;
            grad(jnode, jlev, LATdLAT) -= 2. * avgS(iedge, jlev, LATdLAT) * metric_y;
        }
    }
}

// ================================================================================

void Nabla::divergence(const Field& vector_field, Field& div_field) const {
    const double radius  = fvm_->radius();
    const double deg2rad = M_PI / 180.;

    const mesh::Edges& edges = fvm_->mesh().edges();
    const mesh::Nodes& nodes = fvm_->mesh().nodes();

    const idx_t nnodes = fvm_->node_columns().nb_nodes();
    const idx_t nedges = fvm_->edge_columns().nb_edges();

    const auto vector =
        vector_field.levels()
            ? array::make_view<double, 3>(vector_field).slice(Range::all(), Range::all(), Range::all())
            : array::make_view<double, 2>(vector_field).slice(Range::all(), Range::dummy(), Range::all());
    auto div = div_field.levels() ? array::make_view<double, 2>(div_field).slice(Range::all(), Range::all())
                                  : array::make_view<double, 1>(div_field).slice(Range::all(), Range::dummy());

    const idx_t nlev = vector.shape(1);
    if (div.shape(1) != nlev) {
        throw_AssertionFailed("div_field should have same number of levels", Here());
    }

    const auto lonlat_deg     = array::make_view<double, 2>(nodes.lonlat());
    const auto dual_volumes   = array::make_view<double, 1>(nodes.field("dual_volumes"));
    const auto dual_normals   = array::make_view<double, 2>(edges.field("dual_normals"));
    const auto node2edge_sign = array::make_view<double, 2>(nodes.field("node2edge_sign"));
    const auto edge_flags     = array::make_view<int, 1>(edges.flags());
    auto is_pole_edge         = [&](idx_t e) { return Topology::check(edge_flags(e), Topology::POLE); };

    const mesh::Connectivity& node2edge           = nodes.edge_connectivity();
    const mesh::MultiBlockConnectivity& edge2node = edges.node_connectivity();

    array::ArrayT<double> avgS_arr(nedges, nlev, 2ul);
    array::ArrayView<double, 3> avgS = array::make_view<double, 3>(avgS_arr);

    const double scale = deg2rad * deg2rad * radius;

    enum
    {
        LONdLON = 0,
        LATdLAT = 1
    };
    atlas_omp_parallel {
        atlas_omp_for(idx_t jedge = 0; jedge < nedges; ++jedge) {
            double pbc = 1 - is_pole_edge(jedge);

            idx_t ip1 = edge2node(jedge, 0);
            idx_t ip2 = edge2node(jedge, 1);
            double y1 = lonlat_deg(ip1, LAT) * deg2rad;
            double y2 = lonlat_deg(ip2, LAT) * deg2rad;

            double cosy1, cosy2;
            if (metric_approach_ == 0) {
                cosy1 = std::cos(y1) * pbc;
                cosy2 = std::cos(y2) * pbc;
            }
            else {
                cosy1 = cosy2 = std::cos(0.5 * (y1 + y2)) * pbc;
            }

            double S[2] = {dual_normals(jedge, LON) * deg2rad, dual_normals(jedge, LAT) * deg2rad};
            double avg[2];

            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                double u1 = vector(ip1, jlev, LON);
                double u2 = vector(ip2, jlev, LON);
                double v1 = vector(ip1, jlev, LAT) * cosy1;
                double v2 = vector(ip2, jlev, LAT) * cosy2;

                avg[LON] = (u1 + u2) * 0.5;
                avg[LAT] = (v1 + v2) * 0.5;

                avgS(jedge, jlev, LONdLON) = avg[LON] * S[LON];
                avgS(jedge, jlev, LATdLAT) = avg[LAT] * S[LAT];
            }
        }

        atlas_omp_for(idx_t jnode = 0; jnode < nnodes; ++jnode) {
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                div(jnode, jlev) = 0.;
            }
            for (idx_t jedge = 0; jedge < node2edge.cols(jnode); ++jedge) {
                idx_t iedge = node2edge(jnode, jedge);
                if (iedge < nedges) {
                    double add = node2edge_sign(jnode, jedge);
                    for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                        div(jnode, jlev) += add * (avgS(iedge, jlev, LONdLON) + avgS(iedge, jlev, LATdLAT));
                    }
                }
            }
            const double y = lonlat_deg(jnode, LAT) * deg2rad;
            double metric  = 1. / (dual_volumes(jnode) * scale * std::cos(y));
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                div(jnode, jlev) *= metric;
            }
        }
    }
}

void Nabla::curl(const Field& vector_field, Field& curl_field) const {
    const double radius  = fvm_->radius();
    const double deg2rad = M_PI / 180.;

    const mesh::Edges& edges = fvm_->mesh().edges();
    const mesh::Nodes& nodes = fvm_->mesh().nodes();

    const idx_t nnodes = fvm_->node_columns().nb_nodes();
    const idx_t nedges = fvm_->edge_columns().nb_edges();

    const auto vector =
        vector_field.levels()
            ? array::make_view<double, 3>(vector_field).slice(Range::all(), Range::all(), Range::all())
            : array::make_view<double, 2>(vector_field).slice(Range::all(), Range::dummy(), Range::all());
    auto curl = curl_field.levels() ? array::make_view<double, 2>(curl_field).slice(Range::all(), Range::all())
                                    : array::make_view<double, 1>(curl_field).slice(Range::all(), Range::dummy());

    const idx_t nlev = vector.shape(1);
    if (curl.shape(1) != nlev) {
        throw_AssertionFailed("curl field should have same number of levels", Here());
    }

    const auto lonlat_deg     = array::make_view<double, 2>(nodes.lonlat());
    const auto dual_volumes   = array::make_view<double, 1>(nodes.field("dual_volumes"));
    const auto dual_normals   = array::make_view<double, 2>(edges.field("dual_normals"));
    const auto node2edge_sign = array::make_view<double, 2>(nodes.field("node2edge_sign"));
    const auto edge_flags     = array::make_view<int, 1>(edges.flags());
    auto is_pole_edge         = [&](idx_t e) { return Topology::check(edge_flags(e), Topology::POLE); };


    const mesh::Connectivity& node2edge           = nodes.edge_connectivity();
    const mesh::MultiBlockConnectivity& edge2node = edges.node_connectivity();

    array::ArrayT<double> avgS_arr(nedges, nlev, 2ul);
    array::ArrayView<double, 3> avgS = array::make_view<double, 3>(avgS_arr);

    const double scale = deg2rad * deg2rad * radius;

    enum
    {
        LONdLAT = 0,
        LATdLON = 1
    };


    atlas_omp_parallel {
        atlas_omp_for(idx_t jedge = 0; jedge < nedges; ++jedge) {
            idx_t ip1 = edge2node(jedge, 0);
            idx_t ip2 = edge2node(jedge, 1);
            double y1 = lonlat_deg(ip1, LAT) * deg2rad;
            double y2 = lonlat_deg(ip2, LAT) * deg2rad;

            double pbc = 1 - is_pole_edge(jedge);
            double cosy1;
            double cosy2;
            if (metric_approach_ == 0) {
                cosy1 = std::cos(y1) * pbc;
                cosy2 = std::cos(y2) * pbc;
            }
            else {
                cosy1 = cosy2 = std::cos(0.5 * (y1 + y2)) * pbc;
            }

            double S[2] = {dual_normals(jedge, LON) * deg2rad, dual_normals(jedge, LAT) * deg2rad};
            double avg[2];

            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                double u1 = vector(ip1, jlev, LON) * cosy1;
                double u2 = vector(ip2, jlev, LON) * cosy2;
                double v1 = vector(ip1, jlev, LAT);
                double v2 = vector(ip2, jlev, LAT);

                avg[LON] = (u1 + u2) * 0.5;
                avg[LAT] = (v1 + v2) * 0.5;

                avgS(jedge, jlev, LONdLAT) = avg[LON] * S[LAT];
                avgS(jedge, jlev, LATdLON) = avg[LAT] * S[LON];
            }
        }

        atlas_omp_for(idx_t jnode = 0; jnode < nnodes; ++jnode) {
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                curl(jnode, jlev) = 0.;
            }
            for (idx_t jedge = 0; jedge < node2edge.cols(jnode); ++jedge) {
                idx_t iedge = node2edge(jnode, jedge);
                if (iedge < nedges) {
                    double add = node2edge_sign(jnode, jedge);
                    for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                        curl(jnode, jlev) += add * (avgS(iedge, jlev, LATdLON) - avgS(iedge, jlev, LONdLAT));
                    }
                }
            }
            double y      = lonlat_deg(jnode, LAT) * deg2rad;
            double metric = 1. / (dual_volumes(jnode) * scale * std::cos(y));
            for (idx_t jlev = 0; jlev < nlev; ++jlev) {
                curl(jnode, jlev) *= metric;
            }
        }
    }
}

void Nabla::laplacian(const Field& scalar, Field& lapl) const {
    Field grad(fvm_->node_columns().createField<double>(option::name("grad") | option::levels(scalar.levels()) |
                                                        option::variables(2)));
    gradient(scalar, grad);
    if (fvm_->node_columns().halo().size() < 2) {
        fvm_->node_columns().haloExchange(grad);
    }
    divergence(grad, lapl);
}

const FunctionSpace& Nabla::functionspace() const {
    return fvm_->node_columns();
}

}  // namespace fvm
}  // namespace numerics
}  // namespace atlas
