//#####################################################################
// Copyright 2010.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Level_Sets/EXTRAPOLATION_HIGHER_ORDER.h>
#include <Core/Arrays/INDIRECT_ARRAY.h>
#include <Core/Boundaries/PERIODIC_BOUNDARY_HELPER.h>
#include <Core/Grids/CELL_ITERATOR.h>
#include <Core/Grids/FACE_ITERATOR.h>
#include <Core/Grids/NODE_ITERATOR.h>
#include <Core/Level_Sets/FAST_MARCHING_METHOD_UNIFORM.h>
#include <Core/Level_Sets/LEVELSET.h>
#include <Core/Level_Sets/REINITIALIZATION.h>
#include <Core/Math_Tools/INTERVAL.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
#include <Core/Viewer/VIEWER_OUTPUT.h>
#include <climits>
using namespace PhysBAM;
template <class T> static enable_if_t<IS_NUMBER<T>::value>
set_if_abs_less(T a,T b,T& x,T y,T z){if(abs(a)<abs(b)) x=y;else x=z;}    
static void set_if_abs_less(double a,double b,double& x,double y,double z){if(abs(a)<abs(b)) x=y;else x=z;}
template<class T,int d>
static void set_if_abs_less(const SYMMETRIC_MATRIX<T,d>& a,const SYMMETRIC_MATRIX<T,d>& b,SYMMETRIC_MATRIX<T,d>& x,const SYMMETRIC_MATRIX<T,d>& y,const SYMMETRIC_MATRIX<T,d>& z)
{
    for(int i=0;i<d*(d+1)/2;i++)
        set_if_abs_less(a.x[i],b.x[i],x.x[i],y.x[i],z.x[i]);
}
template<class T,int m,int n>
static void set_if_abs_less(const MATRIX<T,m,n>& a,const MATRIX<T,m,n>& b,MATRIX<T,m,n>& x,const MATRIX<T,m,n>& y,const MATRIX<T,m,n>& z)
{
    for(int i=0;i<m*n;i++)
        set_if_abs_less(a.x[i],b.x[i],x.x[i],y.x[i],z.x[i]);
}
//#####################################################################
// Constructor
//#####################################################################
template<class TV> EXTRAPOLATION_HIGHER_ORDER<TV>::
EXTRAPOLATION_HIGHER_ORDER(const LEVELSET<TV>& phi,const GRID<TV>& grid,
    AT_NODES,std::function<bool(const IV& index)> inside_mask,
    const RANGE<IV>& array_domain,int order,int fill_width,
    const BV& periodic,const IV& periodic_size)
    :phi(phi),grid(grid),order(order),fill_width(fill_width),
    periodic(periodic),periodic_size(periodic_size)
{
    PHYSBAM_ASSERT(order>=1 && order<=3);
    Register_Nodes(grid,array_domain,inside_mask);
}

//#####################################################################
// Constructor
//#####################################################################
template<class TV> EXTRAPOLATION_HIGHER_ORDER<TV>::
EXTRAPOLATION_HIGHER_ORDER(const LEVELSET<TV>& phi,const GRID<TV>& grid,
    AT_NODES,std::function<bool(const IV& index)> inside_mask,
    const RANGE<IV>& array_domain,int order,int fill_width,
    const BV& periodic)
    :EXTRAPOLATION_HIGHER_ORDER(phi,grid,at_nodes,inside_mask,array_domain,
        order,fill_width,periodic,grid.numbers_of_cells)
{
}

//#####################################################################
// Constructor
//#####################################################################
template<class TV> EXTRAPOLATION_HIGHER_ORDER<TV>::
EXTRAPOLATION_HIGHER_ORDER(const LEVELSET<TV>& phi,const GRID<TV>& grid,
    AT_CELLS,std::function<bool(const IV& index)> inside_mask,
    const RANGE<IV>& array_domain,int order,int fill_width,
    const BV& periodic,const IV& periodic_size)
    :EXTRAPOLATION_HIGHER_ORDER(phi,grid.Get_Regular_Grid_At_MAC_Positions(),
        at_nodes,inside_mask,array_domain,order,fill_width,periodic,
        periodic_size)
{
}

//#####################################################################
// Constructor
//#####################################################################
template<class TV> EXTRAPOLATION_HIGHER_ORDER<TV>::
EXTRAPOLATION_HIGHER_ORDER(const LEVELSET<TV>& phi,const GRID<TV>& grid,
    AT_CELLS,std::function<bool(const IV& index)> inside_mask,
    const RANGE<IV>& array_domain,int order,int fill_width,
    const BV& periodic)
    :EXTRAPOLATION_HIGHER_ORDER(phi,grid.Get_Regular_Grid_At_MAC_Positions(),
        at_nodes,inside_mask,array_domain,order,fill_width,periodic,
        grid.numbers_of_cells)
{
}

//#####################################################################
// Function Add_Neighbors
//#####################################################################
template<class TV> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Add_Neighbors(const GRID<TV>& grid,ARRAY<IV>& next,const ARRAY<IV>& neighbors,const IV& index,int unregistered,int registered)
{
    for(int i=0;i<neighbors.m;i++){
        IV neighbor_index=index+neighbors(i);
        Periodic_Index(grid,neighbor_index);
        if(!node_to_index.Valid_Index(neighbor_index)) continue;
        int& n=node_to_index(neighbor_index);
        if(n!=unregistered) continue;
        next.Append(neighbor_index);
        n=registered;}
}
//#####################################################################
// Function Register_Nodes
//#####################################################################
template<class TV> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Register_Nodes(const GRID<TV>& grid,const RANGE<IV>& array_domain,
    std::function<bool(const IV& index)> inside_mask)
{
    node_to_index.Resize(array_domain.Thickened(1),init_all,0); // Need extra ring for the sentinals
    index_to_node.Remove_All();
    index_to_node.Append(IV()+INT_MAX); // First index is the "outside" index.
    normal.Remove_All();
    normal.Append(TV()+FLT_MAX);

    ARRAY<IV> next_inside,next_outside,current;
    ARRAY<IV> neighbors_outside,neighbors_inside;
    for(RANGE_ITERATOR<TV::m> it(RANGE<IV>(IV()-1,IV()+2));it.Valid();it.Next())
        neighbors_outside.Append(it.index);
    for(int d=0;d<TV::m;d++)
        for(int s=-1;s<=1;s+=2){
            IV index;
            index(d)=s;
            neighbors_inside.Append(index);}

    for(RANGE_ITERATOR<TV::m> it(node_to_index.domain,array_domain,RI::ghost);it.Valid();it.Next())
        node_to_index(it.index)=-8;

    for(RANGE_ITERATOR<TV::m> it(array_domain);it.Valid();it.Next())
        node_to_index(it.index)=-1-inside_mask(it.index);
    
    for(FACE_RANGE_ITERATOR<TV::m> it(array_domain,RF::skip_outer);it.Valid();it.Next())
    {
        IV i0=it.face.First_Cell_Index();
        IV i1=it.face.Second_Cell_Index();
        int& n0=node_to_index(i0);
        int& n1=node_to_index(i1);
        if(n0%2==n1%2) continue;
        if(n0==-1){next_outside.Append(i0);n0=-3;}
        else if(n0==-2){next_inside.Append(i0);n0=-4;}
        if(n1==-1){next_outside.Append(i1);n1=-3;}
        else if(n1==-2){next_inside.Append(i1);n1=-4;}
    }

    for(int o=0;o<fill_width;o++){
        current.Exchange(next_outside);
        for(int j=0;j<current.m;j++){
            Register_Index(grid,current(j));
            Add_Neighbors(grid,next_outside,neighbors_outside,current(j),-1,-3);}
        current.Remove_All();}
    next_outside.Remove_All();

    for(int i=1;i<index_to_node.m;i++){
        TV N=phi.Normal(grid.Node_X(index_to_node(i)));
        normal.Append(N);
        for(int d=0;d<TV::m;d++){
            IV index=index_to_node(i);
            int s=N(d)<0?1:-1;
            for(int j=0;j<3;j++){
                Register_Index(grid,index,-3);
                index(d)+=s;}}}

    for(int o=0;o<order;o++){
        solve_indices(o)=INTERVAL<int>(1,index_to_node.m);
        fill_indices(o).min_corner=index_to_node.m;
        current.Exchange(next_inside);
        for(int j=0;j<current.m;j++){
            Register_Index(grid,current(j));
            Add_Neighbors(grid,next_inside,neighbors_inside,current(j),-2,-4);}
        current.Remove_All();}

    for(int o=order-1;o>=0;o--){
        current.Exchange(next_inside);
        for(int j=0;j<current.m;j++){
            Register_Index(grid,current(j));
            Add_Neighbors(grid,next_inside,neighbors_inside,current(j),-2,-4);}
        current.Remove_All();
        fill_indices(o).max_corner=index_to_node.m;}
    next_inside.Remove_All();

    for(int i=normal.m;i<fill_indices(1).max_corner;i++)
        normal.Append(phi.Normal(grid.Node_X(index_to_node(i))));

    VECTOR<STENCIL,TV::m> st;
    stencil.Append(st);
    for(int i=solve_indices(order-1).min_corner;i<solve_indices(order-1).max_corner;i++){
        TV N=normal(i);
        for(int d=0;d<TV::m;d++){
            IV index=index_to_node(i);
            int s=N(d)<0?1:-1;
            for(int j=0;j<3;j++){
                st(d).nodes(j+1)=Lookup_Index(grid,index);
                st(d).scale=s*N(d)*grid.one_over_dX(d);
                index(d)+=s;}
            index(d)-=4*s;
            st(d).nodes(0)=Lookup_Index(grid,index);}
        stencil.Append(st);}
}
//#####################################################################
// Function Fill_un
//#####################################################################
template<class TV> template<class T2> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Fill_un(const GRID<TV>& grid,const TV& one_over_dx,const ARRAY<T2>& x,ARRAY<T2>& xn,int o,int mo) const
{
    xn.Resize(fill_indices(o).max_corner);
    for(int i=fill_indices(o).min_corner;i<fill_indices(o).max_corner;i++){
        const IV& index=index_to_node(i);
        T2 v=T2();
        for(int d=0;d<TV::m;d++){
            IV a=index,b=index;
            a(d)--;
            b(d)++;
            v+=(x(Lookup_Index(grid,b))-x(Lookup_Index(grid,a)))*(T).5*one_over_dx(d)*normal(i)(d);}
        xn(i)=v;}
}
//#####################################################################
// Function Constant_Extrapolate_FE
//#####################################################################
template<class TV> template<class T2> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Extrapolate_FE(const ARRAY<VECTOR<STENCIL,TV::m> >& stencil,const ARRAY<T2>& u,ARRAY<T2>& y,const ARRAY<T2>* z,int o,T dt,T alpha) const
{
    for(int i=solve_indices(o).min_corner;i<solve_indices(o).max_corner;i++){
        T2 dot=T2(),zi=z?(*z)(i):T2(),a=T2();
        for(int d=0;d<TV::m;d++){
            // Second order ENO.
            const STENCIL& s=stencil(i)(d);
            VECTOR<T2,4> f(u(s.nodes(0)),u(s.nodes(1)),u(s.nodes(2)),u(s.nodes(3)));
            VECTOR<T2,3> df(f(1)-f(0),f(2)-f(1),f(3)-f(2));
            VECTOR<T2,2> ddf(df(1)-df(0),df(2)-df(1));
            set_if_abs_less(ddf(0),ddf(1),a,(T).5*(f(2)-f(0)),df(1)-(T).5*ddf(1));
            dot+=a*s.scale;}
        y(i)+=alpha*(u(i)-y(i)-dt*(dot-zi));}
}
//#####################################################################
// Function Constant_Extrapolate_RK2
//#####################################################################
template<class TV> template<class T2> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Extrapolate_RK2(const ARRAY<VECTOR<STENCIL,TV::m> >& stencil,ARRAY<T2>& u,const ARRAY<T2>* z,ARRAY<T2>& tmp,int o,T dt) const
{
    tmp.Resize(u.m);
    for(int i=fill_indices(o).min_corner;i<fill_indices(o).max_corner;i++) tmp(i)=u(i);
    Extrapolate_FE(stencil,u,tmp,z,o,dt,1);
    Extrapolate_FE(stencil,tmp,u,z,o,dt,(T).5);
}
//#####################################################################
// Function Extrapolate
//#####################################################################
template<class TV> template<class T2> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Extrapolate(ARRAYS_ND_BASE<T2,IV>& u,int iterations) const
{
    T dt=grid.dX.Max()/(TV::m+1);
    int num_indices=fill_indices(0).max_corner;
    ARRAY<T2> du[3],tmp(num_indices);
    du[0].Resize(num_indices);
    for(int i=fill_indices(0).min_corner;i<fill_indices(0).max_corner;i++)
        du[0](i)=u(index_to_node(i));
    for(int o=1;o<order;o++)
        Fill_un(grid,grid.one_over_dX,du[o-1],du[o],o,order);
    for(int i=0;i<order;i++)
        du[i](0)=T2()+FLT_MAX/100; // Sentinal values for ENO.
    tmp(0)=T2()+FLT_MAX/100;
    for(int o=order-1;o>=0;o--)
        for(int i=0;i<iterations;i++)
            Extrapolate_RK2(stencil,du[o],(o!=order-1?&du[o+1]:0),tmp,o,dt);
    for(int i=solve_indices(0).min_corner;i<solve_indices(0).max_corner;i++)
        u(index_to_node(i))=du[0](i);
    Fill_Ghost_Periodic(periodic_size,u,periodic);
}
//#####################################################################
// Function Quadratic_Extrapolate
//#####################################################################
// template<class TV> void EXTRAPOLATION_HIGHER_ORDER<TV>::
// Extrapolate_Cell(const GRID<TV>& grid,std::function<bool(const IV& index)> inside_mask,ARRAYS_ND_BASE<T2,IV>& u)
// {
//     Extrapolate_Node(grid.Get_Regular_Grid_At_MAC_Positions(),inside_mask,u);
// }
//#####################################################################
// Function Quadratic_Extrapolate
//#####################################################################
// template<class TV> void EXTRAPOLATION_HIGHER_ORDER<TV>::
// Extrapolate_Face(const GRID<TV>& grid,std::function<bool(const FACE_INDEX<TV::m>& index)> inside_mask,ARRAY<T2,FACE_INDEX<TV::m> >& u)
// {
//     for(int i=0;i<TV::m;i++){
//         if(periodic) combine_ends=VECTOR<bool,TV::m>::Axis_Vector(i);
//         Extrapolate_Node(grid.Get_Face_Grid(i),[&](const IV& index){return inside_mask(FACE_INDEX<TV::m>(i,index));},u.Component(i));}
// }
//#####################################################################
// Function Lookup_Index
//#####################################################################
template<class TV> int EXTRAPOLATION_HIGHER_ORDER<TV>::
Lookup_Index(const GRID<TV>& grid,IV index) const
{
    Periodic_Index(grid,index);
    if(!node_to_index.Valid_Index(index)) return 0;
    int i=node_to_index(index);
    return i>=0?i:0;
}
//#####################################################################
// Function Register_Index
//#####################################################################
template<class TV> int EXTRAPOLATION_HIGHER_ORDER<TV>::
Register_Index(const GRID<TV>& grid,IV index,int only_neg)
{
    Periodic_Index(grid,index);
    if(!node_to_index.Valid_Index(index)) return 0;
    int& n=node_to_index(index);
    if(n<0 && (only_neg>=0 || n==only_neg)) n=index_to_node.Append(index);
    return n<0?0:n;
}
//#####################################################################
// Function Periodic_Index
//#####################################################################
template<class TV> void EXTRAPOLATION_HIGHER_ORDER<TV>::
Periodic_Index(const GRID<TV>& grid,IV& index) const
{
    if(!periodic.Contains(true)) return;
    IV a=index;
    for(int i=0;i<TV::m;i++)
        if(periodic(i)){
            if(index(i)<0) index(i)+=periodic_size(i);
            else if(index(i)>=periodic_size(i)) index(i)-=periodic_size(i);}
}
namespace PhysBAM{
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,1> >;
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,2> >;
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,3> >;
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,1> >;
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,2> >;
template class EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,3> >;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,2> >::Extrapolate<double>(
    ARRAYS_ND_BASE<double,VECTOR<int,2> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,2> >::Extrapolate<float>(
    ARRAYS_ND_BASE<float,VECTOR<int,2> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,1> >::Extrapolate<double>(
    ARRAYS_ND_BASE<double,VECTOR<int,1> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,2> >::Extrapolate<SYMMETRIC_MATRIX<double,2> >(
    ARRAYS_ND_BASE<SYMMETRIC_MATRIX<double,2>,VECTOR<int,2> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,3> >::Extrapolate<SYMMETRIC_MATRIX<double,3> >(
    ARRAYS_ND_BASE<SYMMETRIC_MATRIX<double,3>,VECTOR<int,3> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<double,3> >::Extrapolate<double>(
    ARRAYS_ND_BASE<double,VECTOR<int,3> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,1> >::Extrapolate<float>(
    ARRAYS_ND_BASE<float,VECTOR<int,1> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,2> >::Extrapolate<SYMMETRIC_MATRIX<float,2> >(
    ARRAYS_ND_BASE<SYMMETRIC_MATRIX<float,2>,VECTOR<int,2> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,3> >::Extrapolate<SYMMETRIC_MATRIX<float,3> >(
    ARRAYS_ND_BASE<SYMMETRIC_MATRIX<float,3>,VECTOR<int,3> >&,int) const;
template void EXTRAPOLATION_HIGHER_ORDER<VECTOR<float,3> >::Extrapolate<float>(
    ARRAYS_ND_BASE<float,VECTOR<int,3> >&,int) const;
}
