/*
 * (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. and Interpolation
 */

#pragma once

#include <cmath>
#include <limits>

#include "atlas/array/ArrayView.h"
#include "atlas/functionspace/StructuredColumns.h"
#include "atlas/grid/Stencil.h"
#include "atlas/grid/StencilComputer.h"
#include "atlas/runtime/Exception.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/Point.h"

#include "Cubic3DLimiter.h"
#include "CubicVerticalKernel.h"
#include "LinearHorizontalKernel.h"
#include "QuasiCubicHorizontalKernel.h"

namespace atlas {
namespace interpolation {
namespace method {

struct QuasiCubicLinearPoints {
    constexpr QuasiCubicLinearPoints() {}
    static constexpr std::array<idx_t, 2> j{{1, 2}};
    static constexpr std::array<idx_t, 2> jj{{0, 3}};
    static constexpr std::array<idx_t, 2> jw{{4, 5}};
    static constexpr std::array<idx_t, 2> i{{1, 2}};
    static constexpr std::array<idx_t, 2> ii{{0, 3}};
};

class QuasiCubic3DKernel {
    using Limiter = Cubic3DLimiter;

public:
    QuasiCubic3DKernel(const functionspace::StructuredColumns& fs, const util::Config& config = util::NoConfig()) {
        src_ = fs;
        ATLAS_ASSERT(src_);
        ATLAS_ASSERT(src_.halo() >= 2);
        ATLAS_ASSERT(src_.vertical().size());
        quasi_cubic_horizontal_interpolation_ = QuasiCubicHorizontalKernel(src_, config);
        linear_horizontal_interpolation_      = LinearHorizontalKernel(src_, config);
        vertical_interpolation_               = CubicVerticalKernel(fs.vertical(), config);
        limiter_                              = config.getBool("limiter", false);
    }

private:
    functionspace::StructuredColumns src_;
    QuasiCubicHorizontalKernel quasi_cubic_horizontal_interpolation_;
    LinearHorizontalKernel linear_horizontal_interpolation_;
    CubicVerticalKernel vertical_interpolation_;
    bool limiter_{false};

public:
    static std::string className() { return "QuasiCubic3DKernel"; }
    static constexpr idx_t stencil_width() { return 4; }
    static constexpr idx_t stencil_size() { return 32; }
    static constexpr idx_t stencil_halo() {
        return static_cast<idx_t>(static_cast<double>(stencil_width()) / 2. + 0.5);
    }

public:
    using Stencil = grid::Stencil3D<4>;
    struct Weights {
        std::array<std::array<double, 4>, 4> weights_i;
        std::array<double, 6> weights_j;
        std::array<double, 4> weights_k;
        struct LinearWeights {
            std::array<std::array<double, 2>, 2> weights_i;
            std::array<double, 2> weights_j;
        } linear;
    };


public:
    struct WorkSpace {
        Stencil stencil;
        Weights weights;
    };

    template <typename stencil_t>
    void compute_stencil(const double x, const double y, const double z, stencil_t& stencil) const {
        quasi_cubic_horizontal_interpolation_.compute_stencil(x, y, stencil);
        vertical_interpolation_.compute_stencil(z, stencil);
    }

    template <typename weights_t>
    void compute_weights(const double x, const double y, const double z, weights_t& weights) const {
        Stencil stencil;
        compute_stencil(x, y, z, stencil);
        compute_weights(x, y, z, stencil, weights);
    }


    template <typename stencil_t, typename weights_t>
    void compute_weights(const double x, const double y, const double z, const stencil_t& stencil,
                         weights_t& weights) const {
        quasi_cubic_horizontal_interpolation_.compute_weights(x, y, stencil, weights);


        // Insert more linear weights in available slots (weights_i[0], weights_i[3], weights_j[4], weights_j[5])
        {
            PointXY P1, P2;
            std::array<double, 2> yvec;
            constexpr QuasiCubicLinearPoints pts{};
            // Top and bottom row x-direction
            for (idx_t l = 0; l < 2; ++l) {
                idx_t j         = pts.j[l];   // index in stencil
                idx_t jj        = pts.jj[l];  // row index in weights_i
                auto& weights_i = weights.weights_i[jj];
                src_.compute_xy(stencil.i(pts.i[0], j), stencil.j(j), P1);
                src_.compute_xy(stencil.i(pts.i[1], j), stencil.j(j), P2);
                const double alpha   = (P2.x() - x) / (P2.x() - P1.x());
                weights_i[pts.ii[0]] = alpha;
                weights_i[pts.ii[1]] = 1. - alpha;
                yvec[l]              = P1.y();
            }
            // Compute weights in y-direction
            {
                auto& weights_j    = weights.weights_j;
                const double alpha = (yvec[1] - y) / (yvec[1] - yvec[0]);
                weights_j[4]       = alpha;
                weights_j[5]       = 1. - alpha;
            }
        }

        vertical_interpolation_.compute_weights(z, stencil, weights);
    }

    template <typename stencil_t, typename weights_t, typename array_t>
    typename std::enable_if<(array_t::RANK == 2), typename array_t::value_type>::type interpolate(
        const stencil_t& stencil, const weights_t& weights, const array_t& input) const {
        using Value = typename std::remove_const<typename array_t::value_type>::type;

        std::array<std::array<idx_t, stencil_width()>, stencil_width()> index;
        const auto& wk = weights.weights_k;

        Value output = 0.;

        // Horizontally quasi-cubic part for inner levels ( k = {1,2} )
        {
            // Inner levels, inner rows (cubic in i, cubic in j)   --> 16 points
            const auto& wj = weights.weights_j;
            for (idx_t j = 1; j < 3; ++j) {  // j = {1,2}
                const auto& wi = weights.weights_i[j];
                for (idx_t i = 0; i < 4; ++i) {  // i = {0,1,2,3}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[i] * wj[j];
                    for (idx_t k = 1; k < 3; ++k) {  // k = {1,2}
                        Value w = wij * wk[k];
                        output += w * input(n, stencil.k(k));
                    }
                    index[j][i] = n;
                }
            }
            // Inner levels, outer rows: (linear in i, cubic in j)  --> 8 points
            for (idx_t j = 0; j < 4; j += 3) {  // j = {0,3}
                const auto& wi = weights.weights_i[j];
                for (idx_t i = 1; i < 3; ++i) {  // i = {1,2}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[i] * wj[j];
                    for (idx_t k = 1; k < 3; ++k) {  // k = {1,2}
                        Value w = wij * wk[k];
                        output += w * input(n, stencil.k(k));
                    }
                    index[j][i] = n;
                }
            }
        }
        // Horizontally Linear part for outer levels ( k = {0,3} )
        {
            constexpr QuasiCubicLinearPoints pts{};
            // Outer levels: (linear in i, linear in j) -- > 8 points
            for (idx_t m = 0; m < 2; ++m) {
                idx_t j         = pts.j[m];   // index in stencil ( j = {1,2} )
                idx_t jj        = pts.jj[m];  // row index in weights_i ( jj = {0,3} )
                idx_t jw        = pts.jw[m];  // jw = {4,5};
                const auto& wi  = weights.weights_i[jj];
                const double wj = weights.weights_j[jw];
                for (idx_t l = 0; l < 2; ++l) {
                    idx_t i   = pts.i[l];   // i = {1,2}
                    idx_t ii  = pts.ii[l];  // ii = {0,3}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[ii] * wj;
                    for (idx_t k = 0; k < 4; k += 3) {  // k = {0,3}
                        Value w = wij * wk[k];
                        output += w * input(n, stencil.k(k));
                    }
                }
            }
        }


        if (limiter_) {
            Limiter::limit_scalar(output, index, stencil, input);
        }
        return output;
    }

    template <typename Value>
    struct OutputView1D {
        template <typename Int>
        Value& operator()(Int v) {
            return data_[v];
        }
        template <typename Int>
        Value& operator[](Int v) {
            return data_[v];
        }
        static constexpr int RANK{1};
        OutputView1D(Value* data): data_(data) {}
        using value_type = Value;

        Value* data_;
    };

    template <typename Value>
    OutputView1D<Value> make_outputview(Value* data) const {
        return OutputView1D<Value>(data);
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 3), void>::type interpolate_vars(const stencil_t& stencil,
                                                                                  const weights_t& weights,
                                                                                  const InputArray& input,
                                                                                  OutputArray& output,
                                                                                  const idx_t nvar) const {
        using Value = typename OutputArray::value_type;

        std::array<std::array<idx_t, stencil_width()>, stencil_width()> index;
        const auto& wk = weights.weights_k;

        const Value* _input_;

        for (idx_t v = 0; v < nvar; ++v) {
            output[v] = 0.;
        }

        // Horizontally quasi-cubic part for inner levels ( k = {1,2} )
        {
            // Inner levels, inner rows (cubic in i, cubic in j)   --> 16 points
            const auto& wj = weights.weights_j;
            for (idx_t j = 1; j < 3; ++j) {  // j = {1,2}
                const auto& wi = weights.weights_i[j];
                for (idx_t i = 0; i < 4; ++i) {  // i = {0,1,2,3}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[i] * wj[j];
                    for (idx_t k = 1; k < 3; ++k) {  // k = {1,2}
                        Value w        = wij * wk[k];
                        const idx_t kk = stencil.k(k);
                        _input_        = &(input(n, kk, 0));  // Assumption that input.stride(2) == 1
                        for (idx_t v = 0; v < nvar; ++v) {
                            output[v] += w * _input_[v];
                        }
                    }
                    index[j][i] = n;
                }
            }
            // Inner levels, outer rows: (linear in i, cubic in j)  --> 8 points
            for (idx_t j = 0; j < 4; j += 3) {  // j = {0,3}
                const auto& wi = weights.weights_i[j];
                for (idx_t i = 1; i < 3; ++i) {  // i = {1,2}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[i] * wj[j];
                    for (idx_t k = 1; k < 3; ++k) {  // k = {1,2}
                        Value w        = wij * wk[k];
                        const idx_t kk = stencil.k(k);
                        _input_        = &(input(n, kk, 0));  // Assumption that input.stride(2) == 1
                        for (idx_t v = 0; v < nvar; ++v) {
                            output[v] += w * _input_[v];
                        }
                    }
                    index[j][i] = n;
                }
            }
        }
        // Horizontally Linear part for outer levels ( k = {0,3} )
        {
            constexpr QuasiCubicLinearPoints pts{};
            // Outer levels: (linear in i, linear in j) -- > 8 points
            for (idx_t m = 0; m < 2; ++m) {
                idx_t j         = pts.j[m];   // index in stencil ( j = {1,2} )
                idx_t jj        = pts.jj[m];  // row index in weights_i ( jj = {0,3} )
                idx_t jw        = pts.jw[m];  // jw = {4,5};
                const auto& wi  = weights.weights_i[jj];
                const double wj = weights.weights_j[jw];
                for (idx_t l = 0; l < 2; ++l) {
                    idx_t i   = pts.i[l];   // i = {1,2}
                    idx_t ii  = pts.ii[l];  // ii = {0,3}
                    idx_t n   = src_.index(stencil.i(i, j), stencil.j(j));
                    Value wij = wi[ii] * wj;
                    for (idx_t k = 0; k < 4; k += 3) {  // k = {0,3}
                        Value w        = wij * wk[k];
                        const idx_t kk = stencil.k(k);
                        _input_        = &(input(n, kk, 0));  // Assumption that input.stride(2) == 1
                        for (idx_t v = 0; v < nvar; ++v) {
                            output[v] += w * _input_[v];
                        }
                    }
                }
            }
        }


        if (limiter_) {
            Limiter::limit_vars(index, stencil, input, output, nvar);
        }
    }


    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 2 && OutputArray::RANK == 1), void>::type interpolate(
        const stencil_t& stencil, const weights_t& weights, const InputArray& input, OutputArray& output,
        idx_t r) const {
        output(r) = interpolate(stencil, weights, input);
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 2 && OutputArray::RANK == 2), void>::type interpolate(
        const stencil_t& stencil, const weights_t& weights, const InputArray& input, OutputArray& output, idx_t r,
        idx_t k) const {
        output(r, k) = interpolate(stencil, weights, input);
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 3 && OutputArray::RANK == 3), void>::type interpolate(
        const stencil_t& stencil, const weights_t& weights, const InputArray& input, OutputArray& output, idx_t r,
        idx_t k) const {
        auto output_vars = make_outputview(&output(r, k, 0));
        interpolate_vars(stencil, weights, input, output_vars, output.shape(2));
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 2 && OutputArray::RANK == 3), void>::type interpolate(
        const stencil_t&, const weights_t&, const InputArray&, OutputArray&, idx_t /*r*/, idx_t /*k*/) const {
        ATLAS_NOTIMPLEMENTED;
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 3 && OutputArray::RANK == 1), void>::type interpolate(
        const stencil_t&, const weights_t&, const InputArray&, OutputArray&, idx_t /*r*/) const {
        ATLAS_NOTIMPLEMENTED;
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 3 && OutputArray::RANK == 1), void>::type interpolate(
        const stencil_t&, const weights_t&, const InputArray&, OutputArray&, idx_t /*r*/, idx_t /*k*/) const {
        ATLAS_NOTIMPLEMENTED;
    }

    template <typename stencil_t, typename weights_t, typename InputArray, typename OutputArray>
    typename std::enable_if<(InputArray::RANK == 3 && OutputArray::RANK == 2), void>::type interpolate(
        const stencil_t&, const weights_t&, const InputArray&, OutputArray&, idx_t /*r*/, idx_t /*k*/) const {
        ATLAS_NOTIMPLEMENTED;
    }
};  // namespace method

}  // namespace method
}  // namespace interpolation
}  // namespace atlas
