//#####################################################################
// Copyright 2021, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#ifndef __GENERAL_INTERPOLATION_HELPERS__
#define __GENERAL_INTERPOLATION_HELPERS__

#include <Core/Arrays/ARRAYS_ND_BASE.h>
#include <Core/Arrays/FACE_ARRAYS.h>
#include <Core/Grids/GRID.h>
#include <Core/Interpolation/FACE_LOOKUP_UNIFORM.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
#include <Core/Tensors/TENSOR.h>
#include <Core/Vectors/VECTOR.h>
namespace PhysBAM{

template<class T,class V,int m>
auto Combine(const VECTOR<V,m>& a,const VECTOR<T,m>& b)
{
    V r={};
    for(int i=0;i<m;i++) r+=a(i)*b(i);
    return r;
}

template<class T,class V,int m,int n>
auto Combine(const MATRIX<V,m,n>& a,const VECTOR<T,n>& b)
{
    VECTOR<V,m> r;
    for(int i=0;i<m;i++)
        for(int j=0;j<n;j++)
            r(i)+=a(i,j)*b(j);
    return r;
}

template<class T,class V,int m,int n,int p>
auto Combine(const TENSOR<V,m,n,p>& a,const VECTOR<T,p>& b)
{
    MATRIX<V,m,n> r;
    for(int i=0;i<m;i++)
        for(int j=0;j<n;j++)
            for(int k=0;k<p;k++)
                r(i,j)+=a(i,j,k)*b(k);
    return r;
}

template<class A,class B,int o>
auto Combine_0(const A& a,const VECTOR<B,o>& b)
{
    VECTOR<decltype(Combine(a,b(0))),o> r;
    for(int i=0;i<o;i++) r(i)=Combine(a,b(i));
    return r;
}

template<class A,class B,int o>
auto Combine_1(const VECTOR<A,o>& a,const VECTOR<B,o>& b)
{
    VECTOR<decltype(Combine(a(0),b(0))),o> r;
    for(int i=0;i<o;i++) r(i)=Combine(a(i),b(i));
    return r;
}

template<int a,int s,class T_ARRAY> auto
Get_Data(const T_ARRAY& u,const VECTOR<int,1>& n)
{
    typedef typename T_ARRAY::ELEMENT T;
    VECTOR<T,s+(a==0)> v;
    auto ia=Get_Stencil_Iterator(u,n);
    for(int i=0;i<s+(a==0);i++)
    {
        v(i)=ia.Value();
        ia.Advance(0);
    }
    return v;
}

template<int a,int s,class T_ARRAY> auto
Get_Data(const T_ARRAY& u,const VECTOR<int,2>& n)
{
    typedef typename T_ARRAY::ELEMENT T;
    MATRIX<T,s+(a==0),s+(a==1)> m;
    auto ia=Get_Stencil_Iterator(u,n);
    for(int i=0;i<s+(a==0);i++)
    {
        auto ib=ia;
        for(int j=0;j<s+(a==1);j++)
        {
            m(i,j)=ib.Value();
            ib.Advance(1);
        }
        ia.Advance(0);
    }
    return m;
}

template<int a,int s,class T_ARRAY> auto
Get_Data(const T_ARRAY& u,const VECTOR<int,3>& n)
{
    typedef typename T_ARRAY::ELEMENT T;
    TENSOR<T,s+(a==0),s+(a==1),s+(a==2)> t;
    auto ia=Get_Stencil_Iterator(u,n);
    for(int i=0;i<s+(a==0);i++)
    {
        auto ib=ia;
        for(int j=0;j<s+(a==1);j++)
        {
            auto ic=ib;
            for(int k=0;k<s+(a==2);k++)
            {
                t(i,j,k)=ic.Value();
                ic.Advance(2);
            }
            ib.Advance(1);
        }
        ia.Advance(0);
    }
    return t;
}

template<int axis,class STENCIL,class TV,class IV,class T_ARRAY>
typename T_ARRAY::ELEMENT
Interpolate(const STENCIL& st,const GRID<TV>& grid,const T_ARRAY& u,const std::pair<IV,TV>& p)
{
    auto A=Get_Data<axis,STENCIL::stencil_size>(u,p.first);
    auto r=st.template Value<axis,TV::m-1>(p.second(TV::m-1));
    auto a=Combine_0(A,r);
    if constexpr(TV::m==1) return a.Sum();
    else
    {
        auto s=st.template Value<axis,TV::m-2>(p.second(TV::m-2));
        auto b=Combine_1(a,s);
        if constexpr(TV::m==2) return b.Sum();
        else
        {
            auto t=st.template Value<axis,TV::m-3>(p.second(TV::m-3));
            auto c=Combine_1(b,t);
            return c.Sum();
        }
    }
}

template<int axis,class STENCIL,class TV,class IV,class T_ARRAY>
VECTOR<typename T_ARRAY::ELEMENT,TV::m>
Gradient(const STENCIL& st,const GRID<TV>& grid,const T_ARRAY& u,const std::pair<IV,TV>& p)
{
    typedef VECTOR<typename T_ARRAY::ELEMENT,TV::m> R;
    auto w=grid.one_over_dX;
    auto A=Get_Data<axis,STENCIL::stencil_size>(u,p.first);
    auto rz=st.template Diff<axis,TV::m-1>(p.second(TV::m-1));
    auto az=Combine_0(A,rz);
    if constexpr(TV::m==1) return R(az.Sum()*w.x);
    else
    {
        auto r=st.template Value<axis,TV::m-1>(p.second(TV::m-1));
        auto a=Combine_0(A,r);
        auto s=st.template Value<axis,TV::m-2>(p.second(TV::m-2));
        auto sy=st.template Diff<axis,TV::m-2>(p.second(TV::m-2));
        auto bz=Combine_1(az,s);
        auto by=Combine_1(a,sy);
        if constexpr(TV::m==2) return R(by.Sum()*w.x,bz.Sum()*w.y);
        else
        {
            auto b=Combine_1(a,s);
            auto t=st.template Value<axis,TV::m-3>(p.second(TV::m-3));
            auto tx=st.template Diff<axis,TV::m-3>(p.second(TV::m-3));
            auto cz=Combine_1(bz,t).Sum();
            auto cy=Combine_1(by,t).Sum();
            auto cx=Combine_1(b,tx).Sum();
            return R(cx*w.x,cy*w.y,cz*w.z);
        }
    }
}

template<int axis,class STENCIL,class TV,class IV,class T_ARRAY>
SYMMETRIC_MATRIX<typename T_ARRAY::ELEMENT,TV::m>
Hessian(const STENCIL& st,const GRID<TV>& grid,const T_ARRAY& u,const std::pair<IV,TV>& p)
{
    typedef SYMMETRIC_MATRIX<typename T_ARRAY::ELEMENT,TV::m> R;
    auto w=grid.one_over_dX;
    auto A=Get_Data<axis,STENCIL::stencil_size>(u,p.first);
    auto rzz=st.template Hess<axis,TV::m-1>(p.second(TV::m-1));
    auto azz=Combine_0(A,rzz);
    if constexpr(TV::m==1) return R(azz.Sum()*w.x*w.x);
    else
    {
        auto r=st.template Value<axis,TV::m-1>(p.second(TV::m-1));
        auto rz=st.template Diff<axis,TV::m-1>(p.second(TV::m-1));
        auto a=Combine_0(A,r);
        auto az=Combine_0(A,rz);
        auto s=st.template Value<axis,TV::m-2>(p.second(TV::m-2));
        auto sy=st.template Diff<axis,TV::m-2>(p.second(TV::m-2));
        auto syy=st.template Hess<axis,TV::m-2>(p.second(TV::m-2));
        auto bzz=Combine_1(azz,s);
        auto byz=Combine_1(az,sy);
        auto byy=Combine_1(a,syy);
        if constexpr(TV::m==2) return R(byy.Sum()*w.x*w.x,byz.Sum()*w.x*w.y,bzz.Sum()*w.y*w.y);
        else
        {
            auto b=Combine_1(a,s);
            auto bz=Combine_1(az,s);
            auto by=Combine_1(a,sy);
            auto t=st.template Value<axis,TV::m-3>(p.second(TV::m-3));
            auto tx=st.template Diff<axis,TV::m-3>(p.second(TV::m-3));
            auto txx=st.template Hess<axis,TV::m-3>(p.second(TV::m-3));
            auto czz=Combine_1(bzz,t).Sum();
            auto cyz=Combine_1(byz,t).Sum();
            auto cyy=Combine_1(byy,t).Sum();
            auto cxz=Combine_1(bz,tx).Sum();
            auto cxy=Combine_1(by,tx).Sum();
            auto cxx=Combine_1(b,txx).Sum();
            return R(cxx*w.x*w.x,cxy*w.x*w.y,cxz*w.x*w.z,cyy*w.y*w.y,cyz*w.y*w.z,czz*w.z*w.z);
        }
    }
}

template<class T2>
void Extrema_Update(T2& a0,T2 b0,T2& a1,T2 b1)
{
    a0=Componentwise_Min(a0,b0);
    a1=Componentwise_Max(a1,b1);
}

template<int ss,int axis,class TV,class IV,class T_ARRAY>
VECTOR<typename T_ARRAY::ELEMENT,2>
Extrema(const GRID<TV>& grid,const T_ARRAY& u0,const T_ARRAY& u1,const std::pair<IV,TV>& p)
{
    typedef typename T_ARRAY::ELEMENT T2;
    auto ia0=Get_Stencil_Iterator(u0,p.first);
    auto ia1=Get_Stencil_Iterator(u1,p.first);
    T2 iv0=ia0.Value(),iv1=ia1.Value();
    for(int i=0;i<ss+(axis==0);i++)
    {
        if constexpr(TV::m==1) Extrema_Update(iv0,ia0.Value(),iv1,ia1.Value());
        else
        {
            auto ib0=ia0,ib1=ia1;
            for(int j=0;j<ss+(axis==1);j++)
            {
                if constexpr(TV::m==2) Extrema_Update(iv0,ib0.Value(),iv1,ib1.Value());
                else
                {
                    auto ic0=ib0,ic1=ib1;
                    for(int k=0;k<ss+(axis==2);k++)
                    {
                        Extrema_Update(iv0,ic0.Value(),iv1,ic1.Value());
                        ic0.Advance(2);
                        ic1.Advance(2);
                    }
                }
                ib0.Advance(1);
                ib1.Advance(1);
            }
        }
        ia0.Advance(0);
        ia1.Advance(0);
    }
    return {iv0,iv1};
}

template<class T2,class IV>
struct FLAT_STENCIL_ITERATOR
{
    const IV& stride;
    const T2* p;
    const T2& Value() const {return *p;}
    void Advance(int i) {p+=stride(i);}
};

template<class T2,class IV> FLAT_STENCIL_ITERATOR<T2,IV>
Get_Stencil_Iterator(const ARRAYS_ND_BASE<T2,IV>& u,const IV& index)
{return {u.stride,&u(index)};}

template<class TV,class IV> auto
Get_Stencil_Iterator(const ARRAY<typename TV::SCALAR,FACE_INDEX<TV::m> >& u,const IV& index)
{return Get_Stencil_Iterator(u.V_face,index);}

template<class T_ARRAY,class IV>
struct GENERIC_STENCIL_ITERATOR
{
    const T_ARRAY& array;
    IV index;
    decltype(auto) Value() const {return array(index);}
    void Advance(int i) {index(i)++;}
};

template<class T_ARRAY,class IV> enable_if_t<sizeof(typename T_ARRAY::LOOKUP_TAG)!=0,GENERIC_STENCIL_ITERATOR<T_ARRAY,IV> >
Get_Stencil_Iterator(const T_ARRAY& u,const IV& index)
{
    return {u,index};
}

template<int axis,class T,int m>
auto Diff_Data(const VECTOR<T,m>& a)
{
    VECTOR<T,m-1> r;
    for(int i=0;i<r.m;i++) r(i)=a(i+1)-a(i);
    return r;
}

template<int axis,class T,int m,int n>
auto Diff_Data(const MATRIX<T,m,n>& a)
{
    MATRIX<T,m-(axis==0),n-(axis==1)> r;
    for(int i=0;i<r.m;i++)
        for(int j=0;j<r.n;j++)
            r(i,j)=a(i+(axis==0),j+(axis==1))-a(i,j);
    return r;
}

template<int axis,class T,int m,int n,int p>
auto Diff_Data(const TENSOR<T,m,n,p>& a)
{
    TENSOR<T,m-(axis==0),n-(axis==1),p-(axis==2)> r;
    for(int i=0;i<r.m;i++)
        for(int j=0;j<r.n;j++)
            for(int k=0;k<r.p;k++)
                r(i,j,k)=a(i+(axis==0),j+(axis==1),k+(axis==2))-a(i,j,k);
    return r;
}

template<int axis,class T,int m>
auto Add_Data(const VECTOR<T,m>& a)
{
    VECTOR<T,m-1> r;
    for(int i=0;i<r.m;i++) r(i)=a(i+1)+a(i);
    return r;
}

template<int axis,class T,int m,int n>
auto Add_Data(const MATRIX<T,m,n>& a)
{
    MATRIX<T,m-(axis==0),n-(axis==1)> r;
    for(int i=0;i<r.m;i++)
        for(int j=0;j<r.n;j++)
            r(i,j)=a(i+(axis==0),j+(axis==1))+a(i,j);
    return r;
}

template<int axis,class T,int m,int n,int p>
auto Add_Data(const TENSOR<T,m,n,p>& a)
{
    TENSOR<T,m-(axis==0),n-(axis==1),p-(axis==2)> r;
    for(int i=0;i<r.m;i++)
        for(int j=0;j<r.n;j++)
            for(int k=0;k<r.p;k++)
                r(i,j,k)=a(i+(axis==0),j+(axis==1),k+(axis==2))+a(i,j,k);
    return r;
}
}
#endif
