// -*- tab-width: 4; indent-tabs-mode: nil -*-
#ifndef DUNE_DORIE_UTIL_INTERPOLATOR_HH
#define DUNE_DORIE_UTIL_INTERPOLATOR_HH

#include <cmath>
#include <memory>

#include <dune/common/parametertree.hh>
#include <dune/common/fvector.hh>

namespace Dune {
  namespace Dorie {
    enum InterpolationMethod {nearest, linear, cubic};

    /// Base class for all interpolators
    template<typename T, const int dim>
    class InterpolatorBase
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Domain Domain;
      typedef typename T::Vector Vector;
      typedef typename T::Array Array;
      typedef typename T::Index Index;
      typedef typename T::IndexArray IndexArray;

    protected:
      const ParameterTree& config;
      const Dune::MPIHelper& helper;
      const int verbose;
      const RF eps;

      bool initialized;

      IndexArray N_RF;
      Domain extensions;
      Domain x;

    public:
      const InterpolationMethod method;

      explicit InterpolatorBase(InterpolationMethod method_, const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_):
      config(config_), helper(helper_), verbose(verbose_), eps(1e-9), initialized(false), method(method_)
      {}

      virtual ~InterpolatorBase () = default;

       RF evaluate(const Array& field, const Domain& pos)
      {
        for(Index i = 0; i < dim; i++)
        {
          if(std::abs(pos[i] - x[i]) > eps)
          {
            update_pos(pos);
            x = pos;
            break;
          }
        }
        return field_value(field);
      }

      void initialize(const IndexArray& N_RF_, const Domain& extensions_)
      {
        if(!initialized)
        {
          N_RF = N_RF_;
          extensions = extensions_;
          initialized = true;
        }
        else
          DUNE_THROW(Exception,"Interpolator has already been initialized");
      }

       Domain get_pos() const
      {
        return x;
      }

    protected:

      virtual void update_pos(const Domain& pos) =0;
      virtual RF field_value(const Array& field) const =0;
    };

    template<typename T, const int dim>
    class NearestNeighborInterpolator: public InterpolatorBase<T,dim>
    {
    public:
      explicit NearestNeighborInterpolator() {}
    };

    template<typename T> class NearestNeighborInterpolator<T,2>: public InterpolatorBase<T,2>
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Domain Domain;
      typedef typename T::Array Array;
      typedef typename T::Index Index;
      typedef typename T::IndexArray IndexArray;

      typedef InterpolatorBase<T,2> IB;

      Index ipos;
      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:
      explicit NearestNeighborInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::nearest, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:

       void update_pos(const Domain& pos) override
      {
        Domain unit_pos;
        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;

        Index ipos0 = std::min((Index)round(unit_pos[0]),N_RF[0]-1);
        Index ipos1 = std::min((Index)round(unit_pos[1]),N_RF[1]-1);

        ipos = ipos0 + ipos1*N_RF[0];
      }

      RF field_value(const Array& field) const override
      {
        return field[ipos];
      }
    };

    template<typename T> class NearestNeighborInterpolator<T,3>: public InterpolatorBase<T,3>
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Domain Domain;
      typedef typename T::Array Array;
      typedef typename T::Index Index;
      typedef typename T::IndexArray IndexArray;

      typedef InterpolatorBase<T,3> IB;

      Index ipos;
      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:
      explicit NearestNeighborInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::nearest, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:

      void update_pos(const Domain& pos) override
      {
        Domain unit_pos;
        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;
        unit_pos[2] = std::max(N_RF[2] * pos[2] / extensions[2] - 0.5, 0.0) + eps;

        Index ipos0 = std::min((Index)round(unit_pos[0]),N_RF[0]-1);
        Index ipos1 = std::min((Index)round(unit_pos[1]),N_RF[1]-1);
        Index ipos2 = std::min((Index)round(unit_pos[2]),N_RF[2]-1);

        ipos = ipos0 + ipos1*N_RF[0] + ipos2*N_RF[0]*N_RF[1];
      }

      RF field_value(const Array& field) const override
      {
        return field[ipos];
      }
    };


    template<typename T, const int dim>
    class LinearInterpolator: public InterpolatorBase<T,dim>
    {
    public:
      explicit LinearInterpolator() {}
    };

    template<typename T> class LinearInterpolator<T,2>: public InterpolatorBase<T,2>
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Vector Vector;
      typedef typename T::Domain Domain;
      typedef typename T::Array Array;
      typedef typename T::IndexArray IndexArray;
      typedef typename T::Index Index;
      typedef InterpolatorBase<T,2> IB;

      typedef Dune::FieldVector<IndexArray,2> MultiIndexArray;
      MultiIndexArray ipos;
      Domain dx;

      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:

      explicit LinearInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::linear, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:

      void update_pos(const Domain& pos) override
      {
        Domain unit_pos;
        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;

        Index ipos0_x = std::min((Index)floor(unit_pos[0]),N_RF[0]-1);
        Index ipos0_y = std::min((Index)floor(unit_pos[1]),N_RF[1]-1);
        Index ipos1_x = std::min((Index)ceil(unit_pos[0]),N_RF[0]-1);
        Index ipos1_y = std::min((Index)ceil(unit_pos[1]),N_RF[1]-1);

        ipos[0][0] = ipos0_x + ipos0_y * N_RF[0];
        ipos[0][1] = ipos1_x + ipos0_y * N_RF[0];
        ipos[1][0] = ipos0_x + ipos1_y * N_RF[0];
        ipos[1][1] = ipos1_x + ipos1_y * N_RF[0];

        dx[0] = ipos1_x - unit_pos[0];
        dx[1] = ipos1_y - unit_pos[1];
      }

      RF field_value(const Array& field) const override
      {
      return field[ipos[0][0]] * dx[0] * dx[1]
             + field[ipos[0][1]] * (1-dx[0]) * dx[1]
             + field[ipos[1][0]] * dx[0] * (1-dx[1])
             + field[ipos[1][1]] * (1-dx[0]) * (1-dx[1]);
      }
    };

    // for implementation see https://en.wikipedia.org/wiki/Trilinear_interpolation
    template<typename T> class LinearInterpolator<T,3>: public InterpolatorBase<T,3>
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Vector Vector;
      typedef typename T::Domain Domain;
      typedef typename T::Array Array;
      typedef typename T::Index Index;
      typedef typename T::IndexArray IndexArray;
      typedef InterpolatorBase<T,3> IB;

      typedef Dune::FieldVector<Dune::FieldVector<IndexArray,2>,2> MultiIndexArray;
      MultiIndexArray ipos;
      Domain dx;

      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:
      explicit LinearInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::linear, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:
      void update_pos(const Domain& pos) override
      {
        Domain unit_pos;

        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;
        unit_pos[2] = std::max(N_RF[2] * pos[2] / extensions[2] - 0.5, 0.0) + eps;

        Index ipos0_x = std::min((Index)floor(unit_pos[0]),N_RF[0]-1);
        Index ipos0_y = std::min((Index)floor(unit_pos[1]),N_RF[1]-1);
        Index ipos0_z = std::min((Index)floor(unit_pos[2]),N_RF[2]-1);
        Index ipos1_x = std::min((Index)ceil(unit_pos[0]),N_RF[0]-1);
        Index ipos1_y = std::min((Index)ceil(unit_pos[1]),N_RF[1]-1);
        Index ipos1_z = std::min((Index)ceil(unit_pos[2]),N_RF[2]-1);

        ipos[0][0][0] = ipos0_x + ipos0_y * N_RF[0] + ipos0_z * N_RF[0] * N_RF[1];
        ipos[0][0][1] = ipos1_x + ipos0_y * N_RF[0] + ipos0_z * N_RF[0] * N_RF[1];
        ipos[0][1][0] = ipos0_x + ipos1_y * N_RF[0] + ipos0_z * N_RF[0] * N_RF[1];
        ipos[0][1][1] = ipos1_x + ipos1_y * N_RF[0] + ipos0_z * N_RF[0] * N_RF[1];
        ipos[1][0][0] = ipos0_x + ipos0_y * N_RF[0] + ipos1_z * N_RF[0] * N_RF[1];
        ipos[1][0][1] = ipos1_x + ipos0_y * N_RF[0] + ipos1_z * N_RF[0] * N_RF[1];
        ipos[1][1][0] = ipos0_x + ipos1_y * N_RF[0] + ipos1_z * N_RF[0] * N_RF[1];
        ipos[1][1][1] = ipos1_x + ipos1_y * N_RF[0] + ipos1_z * N_RF[0] * N_RF[1];

        dx[0] = ipos1_x - unit_pos[0];
        dx[1] = ipos1_y - unit_pos[1];
        dx[2] = ipos1_z - unit_pos[2];
      }

      RF field_value(const Array& field) const override
      {
        RF c00 = field[ipos[0][0][0]] * dx[0] + field[ipos[0][0][1]] * (1-dx[0]);
        RF c01 = field[ipos[1][0][0]] * dx[0] + field[ipos[1][0][1]] * (1-dx[0]);
        RF c10 = field[ipos[0][1][0]] * dx[0] + field[ipos[0][1][1]] * (1-dx[0]);
        RF c11 = field[ipos[1][1][0]] * dx[0] + field[ipos[1][1][1]] * (1-dx[0]);

        RF c0 = c00 * dx[1] + c10 * (1-dx[1]);
        RF c1 = c01 * dx[1] + c11 * (1-dx[1]);

        return c0 * dx[2] + c1 * (1-dx[2]);
      }
    };

    template<typename T, const int dim>
    class CubicInterpolator: public InterpolatorBase<T,dim>
    {
    public:
      explicit CubicInterpolator() {}
    };

    template<typename T> class CubicInterpolator<T,2>: public InterpolatorBase<T,2>
    {
    private:
      typedef typename T::RangeField RF;
      typedef typename T::Domain Domain;
      typedef typename T::Vector Vector;
      typedef typename T::Array Array;
      typedef typename T::IndexArray IndexArray;
      typedef typename T::Index Index;
      typedef InterpolatorBase<T,2> IB;

      typedef Dune::FieldVector<Index,4> FourIndexArray;
      typedef Dune::FieldVector<FourIndexArray,4> MultiFourIndexArray;
      MultiFourIndexArray ipos;
      Domain dx;

      typedef Dune::FieldVector<RF,4> FieldVector;
      typedef Dune::FieldVector<FieldVector,4> MultiFieldVector;

      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:

      explicit CubicInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::cubic, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:

      // Implementation inspired by http://www.paulinternet.nl/?page=bicubic
      void update_pos(const Domain& pos) override
      {
        Domain unit_pos;
        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;

        FourIndexArray ipos_x, ipos_y;
        for(Index i=0;i<4;i++)
        {
          ipos_x[i] = std::max((Index)std::min((Index)floor(unit_pos[0]) + i - 1, N_RF[0] - 1),(Index) 0);
          ipos_y[i] = std::max((Index)std::min((Index)floor(unit_pos[1]) + i - 1, N_RF[1] - 1),(Index) 0);
        }

        for(Index i=0;i<4;i++)
        {
          for(Index j=0;j<4;j++)
          {
            ipos[i][j] = ipos_x[i] + ipos_y[j] * N_RF[0];
          }
        }

        dx[0] = unit_pos[0] - ipos_x[1];
        dx[1] = unit_pos[1] - ipos_y[1];
      }

      RF field_value(const Array& field) const override
      {
        MultiFieldVector p1;
        for(Index i=0;i<4;i++)
        {
          for(Index j=0;j<4;j++)
          {
            p1[i][j] = field[ipos[i][j]];
          }
        }

        FieldVector p2;
        for(Index i=0;i<4;i++)
          p2[i] = cubic_interpolate(p1[i], dx[1]);

        return cubic_interpolate(p2, dx[0]);
      }

      RF cubic_interpolate (const FieldVector& p, const RF& x) const {
      	return p[1] + 0.5 * x*(p[2] - p[0] + x*(2.0*p[0] - 5.0*p[1] + 4.0*p[2] - p[3] + x*(3.0*(p[1] - p[2]) + p[3] - p[0])));
      }
    };

    template<typename T> class CubicInterpolator<T,3>: public InterpolatorBase<T,3>
    {
    private:
      typedef typename T::Domain Domain;
      typedef typename T::Array Array;
      typedef typename T::RangeField RF;
      typedef typename T::Index Index;
      typedef typename T::IndexArray IndexArray;
      typedef InterpolatorBase<T,3> IB;

      typedef Dune::FieldVector<Index,4> FourIndexArray;
      typedef Dune::FieldVector<FourIndexArray,4> FourIndexArray2D;
      typedef Dune::FieldVector<FourIndexArray2D,4> FourIndexArray3D;
      FourIndexArray3D ipos;
      Domain dx;

      typedef Dune::FieldVector<RF,4> FieldVector;
      typedef Dune::FieldVector<FieldVector,4> MultiFieldVector2D;
      typedef Dune::FieldVector<MultiFieldVector2D,4> MultiFieldVector3D;

      const IndexArray& N_RF;
      const Domain& extensions;
      const RF& eps;

    public:
      explicit CubicInterpolator(const ParameterTree& config_, const Dune::MPIHelper& helper_, const int verbose_)
      : IB(InterpolationMethod::cubic, config_, helper_, verbose_), N_RF(IB::N_RF), extensions(IB::extensions), eps(IB::eps)
      {
        update_pos(Domain(0.0));
      }

    private:

      // Implementation inspired by http://www.paulinternet.nl/?page=bicubic
      void update_pos(const Domain& pos) override
      {
        Domain unit_pos;
        unit_pos[0] = std::max(N_RF[0] * pos[0] / extensions[0] - 0.5, 0.0) + eps;
        unit_pos[1] = std::max(N_RF[1] * pos[1] / extensions[1] - 0.5, 0.0) + eps;
        unit_pos[2] = std::max(N_RF[2] * pos[2] / extensions[2] - 0.5, 0.0) + eps;

        FourIndexArray ipos_x, ipos_y, ipos_z;
        for(Index i=0;i<4;i++)
        {
          ipos_x[i] = std::max((Index)std::min((Index)floor(unit_pos[0]) + i - 1, N_RF[0] - 1),(Index) 0);
          ipos_y[i] = std::max((Index)std::min((Index)floor(unit_pos[1]) + i - 1, N_RF[1] - 1),(Index) 0);
          ipos_z[i] = std::max((Index)std::min((Index)floor(unit_pos[2]) + i - 1, N_RF[2] - 1),(Index) 0);
        }

        for(Index i=0;i<4;i++)
        {
          for(Index j=0;j<4;j++)
          {
            for(Index k=0;k<4;k++)
            {
              ipos[i][j][k] = ipos_x[i] + ipos_y[j] * N_RF[0] + ipos_z[k] * N_RF[0] * N_RF[1];
            }
          }
        }

        dx[0] = unit_pos[0] - ipos_x[1];
        dx[1] = unit_pos[1] - ipos_y[1];
        dx[2] = unit_pos[2] - ipos_z[1];
      }

      RF field_value(const Array& field) const override
      {
        MultiFieldVector3D p1;
        for(Index i=0;i<4;i++)
        {
          for(Index j=0;j<4;j++)
          {
            for(Index k=0;k<4;k++)
            {
              p1[i][j][k] = field[ipos[i][j][k]];
            }
          }
        }

        FieldVector p2;
        for(Index i=0;i<4;i++)
        {
          p2[i] = bicubic_interpolate(p1[i], dx[1], dx[2]);
        }
        return cubic_interpolate(p2, dx[0]);
      }

      RF bicubic_interpolate (const MultiFieldVector2D& p, const RF& x, const RF& y) const
      {
        FieldVector p1;
        for(Index i=0;i<4;i++)
          p1[i] = cubic_interpolate(p[i], y);
        return cubic_interpolate(p1, x);
      }

      RF cubic_interpolate (const FieldVector& p, const RF& x) const
      {
      	return p[1] + 0.5 * x*(p[2] - p[0] + x*(2.0*p[0] - 5.0*p[1] + 4.0*p[2] - p[3] + x*(3.0*(p[1] - p[2]) + p[3] - p[0])));
      }
    };

    template<typename T, const int dim>
    class InterpolatorFactory
    {
    private:
      typedef InterpolatorBase<T,dim> IB;

    protected:
      const ParameterTree& config;
      const Dune::MPIHelper& helper;
      const int verbose;
      const std::string interpMethod;

    public:
      explicit InterpolatorFactory(const ParameterTree& config_,
                          const Dune::MPIHelper& helper_, const int verbose_)
      : config(config_), helper(helper_), verbose(verbose_)
      , interpMethod(config.get<std::string>("parameters.interpolation"))
      {}

      std::unique_ptr<IB> create()
      {
        std::unique_ptr<IB> interp;
        if(interpMethod == "nearest")
          interp = std::make_unique<NearestNeighborInterpolator<T,dim>>(config, helper, verbose);
        else if (interpMethod == "linear")
          interp = std::make_unique<LinearInterpolator<T,dim>>(config, helper, verbose);
        else if (interpMethod == "cubic")
          interp = std::make_unique<CubicInterpolator<T,dim>>(config, helper, verbose);
        else
          DUNE_THROW(Dune::IOError,"Unrecognized interpolation method encountered: " + interpMethod);

        return interp;
      }
    };
  }
}

#endif