//##################################################################### 
// Copyright 2021, Ritoban Roy-Chowdhury.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Interpolation/DIV_FREE_BALSARA.h>
#include <Core/Interpolation/GENERAL_INTERPOLATION_HELPERS.h>
#include <Core/Interpolation/INTERPOLATION_FACE.h>
#include <Core/Interpolation/INTERPOLATION_HELPERS.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Core/Matrices/MATRIX.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);
}
}

//#####################################################################
// Function Get_Coeffs
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
void DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Get_Coeffs(const GRID<TV>& grid,const T_FACE_ARRAY& u,const IV& center, TV& T,
    TM2& E, TM2& S) const
{
    TV dX=grid.dX;
    DM2 dXi=DM2(grid.one_over_dX);

    VECTOR<TM2, 2> dB;
    VECTOR<TV2, 2> W;
    for (int component=0; component < TV::m; component++) 
    {
        for (int side=0; side < 2; side++) 
        {
            IV i=center.Add_Axis(component, side);
            FACE_INDEX<TV::m> fi(component, i);
            for (int delta_dir=0; delta_dir < TV::m; delta_dir++)
            {
                if (component == delta_dir) continue;
                FACE_INDEX<TV::m> upper(component, i.Add_Axis(delta_dir,1));
                FACE_INDEX<TV::m> lower(component, i.Add_Axis(delta_dir,-1));
                dB(side)(delta_dir, component) = minmod(u(upper) - u(fi), u(fi) - u(lower));
            }
            W(side)(component)=u(fi);
        }
    }


    // ax, by, cz
    TV A=(W(1)-W(0))/dX;

    // first order cross term matrix
    TM2 B=0.5*(dB(1)+dB(0));
    TM2 C=B.Transposed() * dXi;
    // add back in diagonal temrs
    E=C + DM2(A);

    // second order terms
    TM2 P=dXi * (dB(1)-dB(0)) * dXi;
    TV R;
    for (int a=0;a<TV::m;a++) 
    {
        TV b=P.Row(a);
        R(a)=-0.5*b.Sum();
    }
    S=P + DM2(R);

    T=0.5*(W(1)+W(0))-0.25*R*dX*dX;
}
//#####################################################################
// Function Interpolate_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Interpolate_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,
    const TV& X) const -> TV2
{
    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;

    TV T;
    TM2 E, S;
    Get_Coeffs(grid, u, center, T, E, S);

    TV U=E * Y;
    TV V=(S.Transposed() * Y) * Y;

    TV r=T + U + V;
    return r;
}
//#####################################################################
// Function Interpolate_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Interpolate_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,
    int axis) const  -> T2
{
    // I'm not sure I entirely understand what this function is supposed to
    // do, but I'm pretty sure it's just an optimized version of
    // Interpolate_Face that returns just one component, so this is what
    // I'm going to do tentatively.
    return Interpolate_Face(grid, u, X)(axis);
}
//#####################################################################
// Function Extrema_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Extrema_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X) const -> VECTOR<VECTOR<T2,2>,TV::m>
{
    PHYSBAM_NOT_IMPLEMENTED();
}
//#####################################################################
// Function Extrema_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Extrema_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X,int axis) const -> VECTOR<T2,2> 
{
    PHYSBAM_NOT_IMPLEMENTED();
}
//#####################################################################
// Function Gradient_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Gradient_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const  -> VECTOR<TV2,TV::m>
{
    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;

    TV _T;
    TM2 E, S;
    Get_Coeffs(grid, u, center, _T, E, S);
    S=S.Transposed();
    TV D=S*Y;

    VECTOR<TV2,TV::m> r;
    for (int i=0;i<TV::m;i++) {
        for (int k=0;k<TV::m;k++) {
            r(i)(k)=E(i,k);
            // V_{i,k}=S'_{ik} X_i + delta_ik S'_{ij} X_j
            r(i)(k)+=S(i,k)*Y(i);
            if (i==k) {
                r(i)(k)+=D(i);
            }
        }
    }
    return r;
}
//#####################################################################
// Function Gradient_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Gradient_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const -> TV2
{ 
    return Gradient_Face(grid, u, X)(axis);
}
//#####################################################################
// Function Hessian_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Hessian_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const -> VECTOR<SM2,TV::m>
{
    PHYSBAM_NOT_IMPLEMENTED();
}
//#####################################################################
// Function Hessian_Face
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Hessian_Face(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const -> SM2
{
    PHYSBAM_NOT_IMPLEMENTED();
}
//#####################################################################
// Function Interpolate_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Interpolate_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const -> TV2
{
    return Interpolate_Face(grid,u,X);
}
//#####################################################################
// Function Interpolate_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Interpolate_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const -> T2
{
    return Interpolate_Face(grid,u,X,axis);
}
//#####################################################################
// Function Extrema_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Extrema_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X) const -> VECTOR<VECTOR<T2,2>,TV::m>
{
    return Extrema_Face(grid,u0,u1,X);
}
//#####################################################################
// Function Extrema_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Extrema_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u0,const T_FACE_ARRAY& u1,const TV& X,int axis) const -> VECTOR<T2,2>
{
    return Extrema_Face(grid,u0,u1,X,axis);
}
template<class TV,class T2,class T_FACE_ARRAY>
//#####################################################################
// Function Gradient_Face_V
//#####################################################################
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Gradient_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const -> VECTOR<TV2,TV::m>
{
    return Gradient_Face(grid,u,X);
}
//#####################################################################
// Function Gradient_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Gradient_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const -> TV2
{
    return Gradient_Face(grid,u,X,axis);
}
//#####################################################################
// Function Hessian_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Hessian_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X) const -> VECTOR<SM2,TV::m>
{
    return Hessian_Face(grid,u,X);
}
//#####################################################################
// Function Hessian_Face_V
//#####################################################################
template<class TV,class T2,class T_FACE_ARRAY>
auto DIV_FREE_BALSARA<TV,T2,T_FACE_ARRAY>::
Hessian_Face_V(const GRID<TV>& grid,const T_FACE_ARRAY& u,const TV& X,int axis) const -> SM2
{
    return Hessian_Face(grid,u,X,axis);
}
//#####################################################################
template class DIV_FREE_BALSARA<VECTOR<double,2>,double,ARRAY<double,FACE_INDEX<2> > >;
template class DIV_FREE_BALSARA<VECTOR<double,3>,double,ARRAY<double,FACE_INDEX<3> > >;
template class DIV_FREE_BALSARA<VECTOR<float,2>,float,ARRAY<float,FACE_INDEX<2> > >;
template class DIV_FREE_BALSARA<VECTOR<float,3>,float,ARRAY<float,FACE_INDEX<3> > >;
}
