//#####################################################################
// Copyright 2015, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Hybrid_Methods/Iterators/PARTICLE_GRID_WEIGHTS_SPLINE.h>
#include <Core/Grids/GRID.h>
#include <Core/Interpolation/INTERPOLATION_SPLINES.h>
#include <Core/Log/DEBUG_UTILITIES.h>
#include <Core/Math_Tools/cube.h>
#include <Core/Math_Tools/pow.h>
#include <Core/Math_Tools/RANGE_ITERATOR.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Core/Symbolics/SYMBOLIC.h>
namespace PhysBAM{

int binom[7][7]={{1},{1,1},{1,2,1},{1,3,3,1},{1,4,6,4,1},{1,5,10,10,5,1},{1,6,15,20,15,6,1}};

// You can use this to compute the m3 and m4 coefficients.  Normally these are
// rather simple and should be hardcoded for efficiency.
template<int deg,size_t r,size_t s>
std::array<double,r+deg>
Compute_Moment_Coefficients(const std::array<std::array<double,r>,s>& w)
{
    std::array<double,r+deg> c={};
    for(int j=0;j<r;j++)
        for(int i=0;i<s;i++)
        {
            double o=1,p=0.5*s-1-i;
            for(int l=deg;l>=0;l--)
            {
                c[j+l]+=w[i][j]*binom[deg][l]*o;
                o*=p;
            }
        }
    if(deg%2) for(auto&i:c) i=-i;
    return c;
}

//#####################################################################
// Constructor
//#####################################################################
template<class TV,int degree> PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
PARTICLE_GRID_WEIGHTS_SPLINE(const GRID<TV>& grid,int threads)
    :BASE(threads),grid(grid)
{
    T dp=1; // bogus; must be overridden.
    if(degree==2) dp=(T)1/4;
    if(degree==3) dp=(T)1/3;
    if(degree==4) dp=(T)5/12;
    this->stencil_width=degree+1;
    this->Dp=sqr(grid.dX.Max())*dp;
    this->Dp_inv=1/this->Dp;
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV,int degree> PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
~PARTICLE_GRID_WEIGHTS_SPLINE()
{
}
//#####################################################################
// Function Compute
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Compute(const PRECOMPUTE_DATA& pd,typename BASE::SCRATCH& scratch,
    bool want_gradient,bool want_pseudo_gradient) const
{
    int pow_n=pow<TV::m,int>(n);
    scratch.index.Resize(pow_n);
    scratch.weight.Resize(pow_n);
    if(want_gradient) scratch.gradient(0).Resize(pow_n);
    if(want_pseudo_gradient) scratch.gradient(1).Resize(pow_n);

    int k=0;
    for(RANGE_ITERATOR<TV::m> it(RANGE<IV>(IV(),IV()+n));it.Valid();it.Next()){
        scratch.index(k)=pd.base+it.index;
        TV all_w;
        for(int i=0;i<TV::m;i++)
            all_w(i)=pd.w(it.index(i))(i);
        scratch.weight(k)=all_w.Product();
        if(want_gradient){
            TV dw(TV()+1);
            for(int j=0;j<TV::m;j++){
                T dj=pd.dw(it.index(j))(j);
                for(int i=0;i<TV::m;i++)
                    dw(i)*=(i==j)?dj:all_w(j);}
            scratch.gradient(0)(k)=dw;}
        if(want_pseudo_gradient)
        {
            TV Xi=grid.Node_X(scratch.index(k));
            TV dw=scratch.weight(k)*this->Dp_inv*(Xi-pd.X);
            scratch.gradient(1)(k)=dw;
        }
        k++;}
    scratch.m3=pd.m3;
    scratch.m4=pd.m4;
}
//#####################################################################
// Function Compute
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Compute(int p,typename BASE::SCRATCH& scratch,bool want_gradient,
    bool want_pseudo_gradient) const
{
    Compute(precompute_data(p),scratch,want_gradient,want_pseudo_gradient);
}
//#####################################################################
// Function Compute
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Compute(const TV& X,typename BASE::SCRATCH& scratch,bool want_gradient,
    bool want_pseudo_gradient) const
{
    PRECOMPUTE_DATA pd;
    Compute_Precompute_Data(pd,X);
    Compute(pd,scratch,want_gradient,want_pseudo_gradient);
}
//#####################################################################
// Function Compute_Precompute_Data
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Compute_Precompute_Data(PRECOMPUTE_DATA& pd,const TV& X) const
{
    T shift=(T).5*(degree-1);
    pd.base=grid.Cell_I(X-shift*grid.dX);
    pd.X=X;
    TV X_eval=X-grid.Node_X(pd.base);
    TV dx=grid.dX,x=X_eval-dx,x2=x*x;
    TV u=X_eval*grid.one_over_dX-shift; // [0,1]
    auto pw=Make_Powers<4>(u);
    if constexpr(degree==1)
    {
        pd.w=Evaluate_Spline(IS::B1,pw);
        pd.dw=Evaluate_Spline_Diff(IS::B1,pw);
    }
    else if constexpr(degree==2)
    {
        pd.w=Evaluate_Spline(IS::B2,pw);
        pd.dw=Evaluate_Spline_Diff(IS::B2,pw);
        pd.m3=(-(T)1/2*pw(1)+(T)3/2*pw(2)-pw(3))*pow<3>(dx);
        pd.m4=((T)1/16+3*pw(2)-6*pw(3)+3*pw(4))*pow<4>(dx);
    }
    else if constexpr(degree==3)
    {
        pd.w=Evaluate_Spline(IS::B3,pw);
        pd.dw=Evaluate_Spline_Diff(IS::B3,pw);
        pd.m3=TV();
        pd.m4=((T)1/3-pw(2)+2*pw(3)-pw(4))*pow<4>(dx);
    }
    else if constexpr(degree==4)
    {
        pd.w=Evaluate_Spline(IS::B4,pw);
        pd.dw=Evaluate_Spline_Diff(IS::B4,pw);
        pd.m3=TV();
        pd.m4=(T)23/48*pow<4>(dx);
    }
    for(auto&a:pd.dw) a*=grid.one_over_dX;
}
//#####################################################################
// Function Update
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Update(ARRAY_VIEW<const TV> X)
{
    precompute_data.Resize(X.m);
    for(int p=0;p<X.m;p++) Compute_Precompute_Data(precompute_data(p),X(p));
}
//#####################################################################
// Function Particle_Ranges
//#####################################################################
template<class TV,int degree> void PARTICLE_GRID_WEIGHTS_SPLINE<TV,degree>::
Particle_Ranges(ARRAY<RANGE<IV> >& ranges) const
{
    ranges.Resize(precompute_data.m);
    for(int i=0;i<precompute_data.m;i++)
    {
        IV base=precompute_data(i).base;
        ranges(i)={base,base+n};
    }
}
//#####################################################################
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,1>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,2>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,3>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,1>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,2>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,3>,1>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,1>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,2>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,3>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,1>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,2>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,3>,2>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,1>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,2>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,3>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,1>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,2>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,3>,3>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,1>,4>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,2>,4>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<float,3>,4>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,1>,4>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,2>,4>;
template class PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<double,3>,4>;
#ifdef USE_SYMBOLIC
template PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<SYMBOLIC,2>,2>::PARTICLE_GRID_WEIGHTS_SPLINE(GRID<VECTOR<SYMBOLIC,2> > const&,int);
template PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<SYMBOLIC,2>,3>::PARTICLE_GRID_WEIGHTS_SPLINE(GRID<VECTOR<SYMBOLIC,2> > const&,int);
template PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<SYMBOLIC,3>,2>::PARTICLE_GRID_WEIGHTS_SPLINE(GRID<VECTOR<SYMBOLIC,3> > const&,int);
template PARTICLE_GRID_WEIGHTS_SPLINE<VECTOR<SYMBOLIC,3>,3>::PARTICLE_GRID_WEIGHTS_SPLINE(GRID<VECTOR<SYMBOLIC,3> > const&,int);
#endif
}
