/*
 * (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 "CubicHorizontalLimiter.h"

#include "eckit/linalg/Triplet.h"

#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"

namespace atlas {
namespace interpolation {
namespace method {

class CubicHorizontalKernel {
    using Limiter  = CubicHorizontalLimiter;
    using Triplet  = eckit::linalg::Triplet;
    using Triplets = std::vector<Triplet>;

public:
    CubicHorizontalKernel() = default;

    CubicHorizontalKernel( const functionspace::StructuredColumns& fs, const util::Config& config = util::NoConfig() ) {
        src_ = fs;
        ATLAS_ASSERT( src_ );
        ATLAS_ASSERT( src_.halo() >= 2 );
        compute_horizontal_stencil_ = grid::ComputeHorizontalStencil( src_.grid(), stencil_width() );
        limiter_                    = config.getBool( "limiter", false );
    }

protected:
    functionspace::StructuredColumns src_;
    grid::ComputeHorizontalStencil compute_horizontal_stencil_;
    bool limiter_{false};

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

public:
    using Stencil = grid::HorizontalStencil<4>;
    struct Weights {
        std::array<std::array<double, 4>, 4> weights_i;
        std::array<double, 4> weights_j;
    };

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

    template <typename stencil_t>
    void compute_stencil( const double x, const double y, stencil_t& stencil ) const {
        compute_horizontal_stencil_( x, y, stencil );
    }

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


    template <typename stencil_t, typename weights_t>
    void compute_weights( const double x, const double y, const stencil_t& stencil, weights_t& weights ) const {
        PointXY P1, P2;
        std::array<double, 4> yvec;
        for ( idx_t j = 0; j < stencil_width(); ++j ) {
            auto& weights_i = weights.weights_i[j];
            src_.compute_xy( stencil.i( 1, j ), stencil.j( j ), P1 );
            src_.compute_xy( stencil.i( 2, j ), stencil.j( j ), P2 );
            const double alpha               = ( P2.x() - x ) / ( P2.x() - P1.x() );
            const double alpha_sqr           = alpha * alpha;
            const double two_minus_alpha     = 2. - alpha;
            const double one_minus_alpha_sqr = 1. - alpha_sqr;
            weights_i[0]                     = -alpha * one_minus_alpha_sqr / 6.;
            weights_i[1]                     = 0.5 * alpha * ( 1. + alpha ) * two_minus_alpha;
            weights_i[2]                     = 0.5 * one_minus_alpha_sqr * two_minus_alpha;
            weights_i[3]                     = 1. - weights_i[0] - weights_i[1] - weights_i[2];
            yvec[j]                          = P1.y();
        }
        const double dl12 = yvec[0] - yvec[1];
        const double dl13 = yvec[0] - yvec[2];
        const double dl14 = yvec[0] - yvec[3];
        const double dl23 = yvec[1] - yvec[2];
        const double dl24 = yvec[1] - yvec[3];
        const double dl34 = yvec[2] - yvec[3];
        const double dcl1 = dl12 * dl13 * dl14;
        const double dcl2 = -dl12 * dl23 * dl24;
        const double dcl3 = dl13 * dl23 * dl34;

        const double dl1 = y - yvec[0];
        const double dl2 = y - yvec[1];
        const double dl3 = y - yvec[2];
        const double dl4 = y - yvec[3];

        auto& weights_j = weights.weights_j;
        weights_j[0]    = ( dl2 * dl3 * dl4 ) / dcl1;
#if defined( _CRAYC ) && ATLAS_BUILD_TYPE_RELEASE
        // prevents FE_INVALID somehow (tested with Cray 8.7)
        ATLAS_ASSERT( !std::isnan( weights_j[0] ) );
#endif
        weights_j[1] = ( dl1 * dl3 * dl4 ) / dcl2;
        weights_j[2] = ( dl1 * dl2 * dl4 ) / dcl3;
        weights_j[3] = 1. - weights_j[0] - weights_j[1] - weights_j[2];
    }

    template <typename stencil_t, typename weights_t, typename array_t>
    typename array_t::value_type interpolate( const stencil_t& stencil, const weights_t& weights,
                                              const array_t& input ) const {
        using Value = typename array_t::value_type;

        std::array<std::array<idx_t, stencil_width()>, stencil_width()> index;
        const auto& weights_j = weights.weights_j;
        Value output          = 0.;
        for ( idx_t j = 0; j < stencil_width(); ++j ) {
            const auto& weights_i = weights.weights_i[j];
            for ( idx_t i = 0; i < stencil_width(); ++i ) {
                idx_t n = src_.index( stencil.i( i, j ), stencil.j( j ) );
                Value w = weights_i[i] * weights_j[j];
                output += w * input[n];
                index[j][i] = n;
            }
        }

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


    template <typename stencil_t, typename weights_t, typename Value, int Rank>
    typename std::enable_if<( Rank == 1 ), void>::type interpolate( const stencil_t& stencil, const weights_t& weights,
                                                                    const array::ArrayView<const Value, Rank>& input,
                                                                    array::ArrayView<Value, Rank>& output,
                                                                    idx_t r ) const {
        std::array<std::array<idx_t, stencil_width()>, stencil_width()> index;
        const auto& weights_j = weights.weights_j;
        output( r )           = 0.;
        for ( idx_t j = 0; j < stencil_width(); ++j ) {
            const auto& weights_i = weights.weights_i[j];
            for ( idx_t i = 0; i < stencil_width(); ++i ) {
                idx_t n = src_.index( stencil.i( i, j ), stencil.j( j ) );
                Value w = static_cast<Value>( weights_i[i] * weights_j[j] );
                output( r ) += w * input[n];
                index[j][i] = n;
            }
        }

        if ( limiter_ ) {
            Limiter::limit( index, input, output, r );
        }
    }


    template <typename stencil_t, typename weights_t, typename Value, int Rank>
    typename std::enable_if<( Rank == 2 ), void>::type interpolate( const stencil_t& stencil, const weights_t& weights,
                                                                    const array::ArrayView<const Value, Rank>& input,
                                                                    array::ArrayView<Value, Rank>& output,
                                                                    idx_t r ) const {
        std::array<std::array<idx_t, stencil_width()>, stencil_width()> index;
        const auto& weights_j = weights.weights_j;
        const idx_t Nk        = output.shape( 1 );
        for ( idx_t k = 0; k < Nk; ++k ) {
            output( r, k ) = 0.;
        }
        for ( idx_t j = 0; j < stencil_width(); ++j ) {
            const auto& weights_i = weights.weights_i[j];
            for ( idx_t i = 0; i < stencil_width(); ++i ) {
                idx_t n = src_.index( stencil.i( i, j ), stencil.j( j ) );
                Value w = static_cast<Value>( weights_i[i] * weights_j[j] );
                for ( idx_t k = 0; k < Nk; ++k ) {
                    output( r, k ) += w * input( n, k );
                }
                index[j][i] = n;
            }
        }

        if ( limiter_ ) {
            Limiter::limit( index, input, output, r );
        }
    }

    template <typename array_t>
    typename array_t::value_type operator()( const double x, const double y, const array_t& input ) const {
        Stencil stencil;
        compute_stencil( x, y, stencil );
        Weights weights;
        compute_weights( x, y, stencil, weights );
        return interpolate( stencil, weights, input );
    }

    template <typename array_t>
    typename array_t::value_type interpolate( const PointLonLat& p, const array_t& input, WorkSpace& ws ) const {
        compute_stencil( p.lon(), p.lat(), ws.stencil );
        compute_weights( p.lon(), p.lat(), ws.stencil, ws.weights );
        return interpolate( ws.stencil, ws.weights, input );
    }

    // Thread private workspace
    Triplets compute_triplets( const idx_t row, const double x, const double y, WorkSpace& ws ) const {
        Triplets triplets;
        triplets.reserve( stencil_size() );
        insert_triplets( row, x, y, triplets, ws );
        return triplets;
    }

    Triplets reserve_triplets( size_t N ) {
        Triplets triplets;
        triplets.reserve( N * stencil_size() );
        return triplets;
    }

    Triplets allocate_triplets( size_t N ) { return Triplets( N * stencil_size() ); }

    void insert_triplets( const idx_t row, const PointXY& p, Triplets& triplets, WorkSpace& ws ) const {
        insert_triplets( row, p.x(), p.y(), triplets, ws );
    }

    void insert_triplets( const idx_t row, const double x, const double y, Triplets& triplets, WorkSpace& ws ) const {
        compute_horizontal_stencil_( x, y, ws.stencil );
        compute_weights( x, y, ws.stencil, ws.weights );
        const auto& wj = ws.weights.weights_j;

        idx_t pos = row * stencil_size();
        for ( idx_t j = 0; j < stencil_width(); ++j ) {
            const auto& wi = ws.weights.weights_i[j];
            for ( idx_t i = 0; i < stencil_width(); ++i ) {
                idx_t col       = src_.index( ws.stencil.i( i, j ), ws.stencil.j( j ) );
                double w        = wi[i] * wj[j];
                triplets[pos++] = Triplet( row, col, w );
            }
        }
    }
};

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