//##################################################################### 
// Copyright 2022, Ritoban Roy-Chowdhury.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#ifndef __DIV_FREE_BALSARA_WENO__
#define __DIV_FREE_BALSARA_WENO__

#include <Core/Grids/FACE_INDEX.h>
#include <Core/Interpolation/DIV_FREE_BALSARA.h>
#include <Core/Interpolation/INTERPOLATION_FACE.h>
#include <Core/Interpolation/INTERPOLATION_HELPERS.h>
#include <Core/Vectors/VECTOR.h>

namespace PhysBAM 
{
namespace
{
template<class T>
static T minmod(T a, T b)
{
    if(a * b < 0) return 0;
    return minmag(a, b);
}
}

template<class TV, class T2, int order>
class DIV_FREE_BALSARA_WENO:public INTERPOLATION_FACE<TV, T2>
{
    // Implements the high order WENO-based interpolation schemes described in 
    // Balsara, D. S. (2009). Divergence-free reconstruction of magnetic fields and WENO schemes for magnetohydrodynamics. Journal of Computational Physics, 228(14), 5040-5056.
    //
    // Note that paper is a companion paper to the following more general paper on high-order ADER-WENO (Arbitrary Derivative Rienmann Problem, Weighted Essentially Non-Oscillatory) schemes for magnetohydrodynamics. In particular, this paper describes the actual WENO limiting procedure that is used for finding the terms like B_yy^x+ in the previous paper.
    // Balsara, D. S., Rumpf, T., Dumbser, M., & Munz, C. D. (2009). Efficient, high accuracy ADER-WENO schemes for hydrodynamics and divergence-free magnetohydrodynamics. Journal of Computational Physics, 228(7), 2480-2516.
    //
    // Finally, for the formulas used for weighting the various schemes, the above Balsara papers adopt the weights used in the following 1996 paper (with some exceptions where they just let the central stencils have a weight 100x the side stencils)
    // Jiang, G. S., & Shu, C. W. (1996). Efficient implementation of weighted ENO schemes. Journal of computational physics, 126(1), 202-228.

    typedef typename TV::SCALAR T;typedef VECTOR<int, TV::m> IV;
    typedef VECTOR<T2,TV::m> TV2;typedef SYMMETRIC_MATRIX<T2,TV::m> SM2;
    typedef MATRIX<T2, TV::m> TM2;typedef DIAGONAL_MATRIX<T2,TV::m> DM2;
    typedef INTERPOLATION_FACE<TV, T2> BASE;
    typedef ARRAY<T2, FACE_INDEX<TV::m>> T_FACE_ARRAY ;

    // TODO: figure out what to do with the stencil_size
    constexpr static int stencil_size=2*order - 1;
    constexpr static int const_extrap=false;

public:
    DIV_FREE_BALSARA_WENO() {};

    T WENO_Limiter(const FACE_INDEX<TV::m>& fi, const IV& subscript, const T_FACE_ARRAY& vel) const
    {
        if (subscript==IV())
        {
            return vel(fi);
        }

        int num_nonzeros=0;
        int direction;
        for(int a=0;a<TV::m;a++) 
        {
            if(subscript(a)!=0) 
            {
                direction=a;
                num_nonzeros++;
            }
        }

        if (num_nonzeros == 1)
        {
            if constexpr(order == 2) {
                // Perform 1D WENO limiting
                FACE_INDEX<TV::m> lower(fi.axis, fi.index.Add_Axis(direction, -1));
                FACE_INDEX<TV::m> upper(fi.axis, fi.index.Add_Axis(direction,  1));
                //return vel(upper) - vel(lower);
                return minmod(vel(upper)-vel(fi), vel(fi)-vel(lower));
            } else if constexpr(order == 3) {
                auto u=[&vel,&fi,&direction](int i)
                {
                    IV idx=fi.index.Add_Axis(direction, i);
                    return vel({fi.axis, idx});
                };
                T left_ux=-2*u(-1) + 0.5*u(-2) + 1.5*u(0);
                T left_uxx=0.5*(u(-2) - 2*u(-1) + u(0));

                T center_ux=0.5*(u(1)-u(-1));
                T center_uxx=0.5*(u(-1) - 2*u(0) + u(1));

                T right_ux=-1.5*u(0)+2*u(1)-0.5*u(2);
                T right_uxx=0.5*(u(0)-2*u(1)+u(2));

                T left_IS=left_ux*left_ux + (T)13/3*left_uxx*left_uxx;
                T center_IS=center_ux*center_ux + (T)13/3*center_uxx*center_uxx;
                T right_IS=right_ux*right_ux + (T)13/3*right_uxx*right_uxx;

                VECTOR<T, order> IS(left_IS, center_IS, right_IS);
                // the paper gives the central stencil a linear weight 100x the outer stencils
                VECTOR<T, order> linear_weight(0., 1., 0.);

                T eps=1e-6;
                // Balsara also raises the smoothness measure to the fourth power
                IS += eps;
                VECTOR<T, order> IS4=IS*IS*IS*IS;
                VECTOR<T, order> alpha=linear_weight / IS4;
                alpha /= alpha.Sum();

                T ux = alpha(0)*left_ux + alpha(1)*center_ux + alpha(2)*right_ux;
                T uxx = alpha(0)*left_uxx + alpha(1)*center_uxx + alpha(2)*right_uxx;
                if(subscript(direction)==1) return ux;
                else if (subscript(direction)==2) return uxx;
                else PHYSBAM_FATAL_ERROR("3-rd WENO limiting subscript must be 2 or 2");
            }
        }
        else
        {
            // should be a cross term, something like u_xy
            // 
            // the four possible evaluations are
            //   u( 1,  1) - u(0, 0) - ux - uy - uxx - uyy
            //  -u( 1, -1) + u(0, 0) + ux - uy + uxx + uyy
            //  -u(-1,  1) + u(0, 0) - ux + uy + uxx + uyy
            //   u(-1, -1) - u(0, 0) + ux + uy - uxx - uyy
            // overall,
            //  i and j are 1 and -1
            //  (i*j) = 1 if i and j are the same
            //  (i*j) = -1 if i and j are different
            //
            // each evaluation is
            //   (i*j) * u(i, j) - (i*j) u(0, 0) - (i*j)uxx - (i*j)uyy - j*ux - i*uy

            VECTOR<VECTOR<T, 2>, 2> limits; 
            VECTOR<IV, 2> limit_axes; 
            int loc=0;
            for(int a=0;a<TV::m;a++) {
                if (subscript(a) == 0) continue;
                if (loc >= 2) PHYSBAM_FATAL_ERROR("WENO Limiting for Cross term only supports limiting in two directions, i.e. u_xy");

                T ux = WENO_Limiter(fi, IV::Axis_Vector(a), vel);
                T uxx = WENO_Limiter(fi, 2*IV::Axis_Vector(a), vel);
                limits(loc) = {ux, uxx};
                limit_axes(loc) = IV::Axis_Vector(a);
                loc++;
            }

            VECTOR<T, 4> poss_uxy;
            VECTOR<T, 4> alphas;
            loc = 0;
            for(int i = -1; i <= 1; i+=2) {
                for(int j = -1; j <= 1; j+=2) {
                    T u00 = vel(fi);
                    IV offset = limit_axes(0)*i + limit_axes(1)*j;
                    T uij = vel({fi.axis, fi.index + offset});
                    T uxy = (i*j) * (uij - u00);
                    uxy += -j*limits(0)(0) - i*limits(1)(0)
                      -(i*j)*(limits(0)(1) + limits(1)(1));

                    T eps=1e-6;
                    T IS = 4*pow(limits(0)(1),2) + 4*pow(limits(1)(1),2) + uxy*uxy;
                    IS += eps;
                    T IS4 = IS*IS*IS*IS;
                    // i'm assuming the linear weights are the same for all four stencils
                    // no reason why they shouldn't be, they are symmetric
                    T alpha = 0.25 / IS4;
                    poss_uxy(loc) = uxy;
                    alphas(loc) = alpha;
                    loc++;
                }
            }
            alphas /= alphas.Sum();
            T weighted_uxy = alphas.Dot(poss_uxy);
            return weighted_uxy;
        }
    }


    T Legendre_Polynomial(int n, T x) const
    {
        switch(n)
        {
            case 0:
                return 1;
            case 1:
                return x;
            case 2:
                return x*x - (T)1/12;
            case 3:
                return x*x*x - (T)3/20*x;
            case 4: {
                T x2=x*x;
                return x2*x2 - (T)3/14*x2 + (T)3/560; }
            default:
                PHYSBAM_FATAL_ERROR("DIV_FREE_BALSARA_WENO: Legendre Polynomials only implemented up to degree 4.");
                break;
        }
    }

    T Get_Coeff(int axis, IV subscript, IV cell, const T_FACE_ARRAY& u) const
    {
        if (subscript(axis)==0) 
        {
            FACE_INDEX<TV::m> left(axis, cell);
            FACE_INDEX<TV::m> right(axis, cell.Add_Axis(axis,1));
            T r=0.5*(WENO_Limiter(left, subscript, u) + WENO_Limiter(right, subscript, u));

            if (subscript.Sum()<=1) {
                r -= (T)1./6*Get_Coeff(axis, 2*IV::Axis_Vector(axis)+subscript, cell, u);
            }

            /*
            if constexpr(order >= 4)
                r -= (T)1/70*Get_Coeff(axis, IV::Axis_Vector(axis, 4), cell, u);
            */
            return r;
        }
        else if (subscript(axis)==1)
        {
            FACE_INDEX<TV::m> left(axis, cell);
            FACE_INDEX<TV::m> right(axis, cell.Add_Axis(axis,1));
            subscript(axis)=0;
            T r=WENO_Limiter(right, subscript, u) - WENO_Limiter(left, subscript, u);
            if constexpr(order == 3)
            {
                if (subscript==IV())
                    r -= (T)1./10*Get_Coeff(axis, 3*IV::Axis_Vector(axis) + subscript, cell, u);
            }
            return r;
        } else if (subscript(axis)==2)
        {
            if (subscript==2*IV::Axis_Vector(axis))
            {
                // axx coefficient
                T r=0;
                for(int a=0;a<TV::m;a++)
                {
                    if (a==axis) continue;
                    IV new_subscript = IV::Axis_Vector(axis) + IV::Axis_Vector(a);
                    r -= (T)0.5 * Get_Coeff(a, new_subscript, cell, u);
                }
                return r;
            }
            else
            {
                // axxy axxz
                if(TV::m == 2) {
                    return (T)0;
                }
                IV new_subscript=IV::Constant_Vector(1);
                subscript-=2*IV::Axis_Vector(axis);
                IV av = IV::Constant_Vector(1) - subscript - IV::Axis_Vector(axis);
                int a=0;
                for(;a<TV::m;a++)
                    if (av==IV::Axis_Vector(a))
                        break;
                return -0.25*Get_Coeff(a, new_subscript, cell, u);
            }
        } else if (subscript(axis) == 3)
        {
            T r=0;
            for(int a=0;a<TV::m;a++)
            {
                if (a==axis) continue;
                IV new_subscript = 2*IV::Axis_Vector(axis) + IV::Axis_Vector(a);
                r -= (T)1/3 * Get_Coeff(a, new_subscript, cell, u);
            }
            return r;
        } else
        {
            PHYSBAM_FATAL_ERROR("DIV_FREE_BALSARA_WENO: Only second order supported for now.");
        }
    }

    T Get_Term(int axis, IV subscript, TV X, IV cell, const T_FACE_ARRAY& u) const
    {
        T prod=1;
        for(int a=0; a<TV::m;a++)
            prod *= Legendre_Polynomial(subscript(a), X(a));
        T coeff=Get_Coeff(axis, subscript, cell, u);
        //LOG::printf("%p_(%p) = %p\n", axis, subscript, coeff);
        return coeff*prod;
    }

    TV2 Interpolate_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const
    {
        auto p = INTERPOLATION_HELPERS::Clamped_Index<stencil_size,const_extrap>(grid, at_cells, u.Domain_Indices(), X);
        IV center = p.first + 1;
        TV center_X = grid.Cell_X(center);
        TV Y = (X - center_X)*grid.one_over_dX;
        //LOG::printf("center: %P at %P\n", center, center_X);

        TV r;

        for (int a=0;a<TV::m;a++) 
        {
            auto add_term=[this,a,&Y,&center,&u,&r](IV subscript)
            {
                T term=Get_Term(a, subscript, Y, center, u);
                r(a) += term;
            };
            if constexpr(order >= 1)
            {
                add_term(IV()); // a0
                for (int s=0;s<TV::m;s++) 
                    add_term(IV::Axis_Vector(s)); // ax, ay, az
            }
            if constexpr(order >= 2)
            {
                // axx, axy, axz
                IV prefix=IV::Axis_Vector(a);
                for (int s=0;s<TV::m;s++)
                    add_term(prefix.Add_Axis(s, 1));
            }
            if constexpr(order >= 3)
            {
                // yy xyy,
                // yz, xyz
                // zz, xzz,
                IV prefix=IV::Axis_Vector(a);
                if constexpr(TV::m == 2)
                {
                    IV subscript=2*IV::Axis_Vector(1-a);
                    add_term(subscript);
                    add_term(prefix+subscript);
                }
                if constexpr(TV::m == 3)
                {
                    for(int i=0;i<=2;i++)
                    {
                        IV a1=IV::Axis_Vector((a + 2) % 3);
                        IV a2=IV::Axis_Vector((a + 1) % 3);
                        IV subscript=i*a1 + (2-i)*a2;
                        add_term(subscript);
                        add_term(prefix+subscript);
                    }
                }

                // axxx, axxy, axxz
                prefix=2*IV::Axis_Vector(a);
                for (int s=0;s<TV::m;s++)
                    add_term(prefix.Add_Axis(s, 1));
            }
            /*if constexpr(order >= 4)
            {
                // ayyy axyyy
                // ayyz axyyz
                // ayzz axyzz
                // azzz axzzz
                IV prefix=IV::Axis_Vector(a, 1);
                // TODO!
                // (0, 3, 0)
                // (0, 2, 1)
                // (0, 1, 2)
                // (0, 3, 0)

                // axxyy axxzz
                prefix=IV::Axis_Vector(a, 2);
                for (int s=0;s<TV::m;s++)
                    if (s != a)
                        add_term(prefix.Add_Axis(s, 2));

                // axxxx axxxy axxxz
                prefix=IV::Axis_Vector(a, 3);
                for (int s=0;s<TV::m;s++)
                    add_term(prefix.Add_Axis(s, 1));
            } */
        }

        return r;
    }


    T2 Interpolate_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const 
    {
        return Interpolate_Face(grid, u, X)(axis);
    }

    VECTOR<VECTOR<T2,2>,TV::m> Extrema_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X) const
    { PHYSBAM_NOT_IMPLEMENTED(); }

    VECTOR<T2,2> Extrema_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X,int axis) const
    { PHYSBAM_NOT_IMPLEMENTED(); }

    VECTOR<TV2,TV::m> Gradient_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const 
    {  PHYSBAM_NOT_IMPLEMENTED(); }

    TV2 Gradient_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const
    { return Gradient_Face(grid, u, X)(axis); }

    VECTOR<SM2,TV::m> Hessian_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const
    { PHYSBAM_NOT_IMPLEMENTED(); }

    SM2 Hessian_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const
    { PHYSBAM_NOT_IMPLEMENTED(); }


protected:
    virtual TV2 Interpolate_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const
    {return Interpolate_Face(grid,u,X);}
    virtual T2 Interpolate_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const
    {return Interpolate_Face(grid,u,X,axis);}
    virtual VECTOR<VECTOR<T2,2>,TV::m> Extrema_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X) const
    {return Extrema_Face(grid,u0,u1,X);}
    virtual VECTOR<T2,2> Extrema_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X,int axis) const
    {return Extrema_Face(grid,u0,u1,X,axis);}
    virtual VECTOR<TV2,TV::m> Gradient_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const
    {return Gradient_Face(grid,u,X);}
    virtual TV2 Gradient_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const
    {return Gradient_Face(grid,u,X,axis);}
    virtual VECTOR<SM2,TV::m> Hessian_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const
    {return Hessian_Face(grid,u,X);}
    virtual SM2 Hessian_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const
    {return Hessian_Face(grid,u,X,axis);}
//#####################################################################
};

}

#endif
