#include <Core/Poisson/COMPUTE_LAPLACE_MATRIX.h>
#include <Core/Arrays/FLOOD_FILL.h>
#include <Core/Math_Tools/RANGE_ITERATOR.h>
namespace PhysBAM{
//#####################################################################
// Function Compute_Laplace_Matrix
//#####################################################################
template<class TV,class T>
void Compute_Laplace_Matrix(LAPLACE_MATRIX<TV>& lm,ARRAY<ARRAY<T>>& rhs,const ARRAYS_ND_BASE<T,VECTOR<int,TV::m>>& f,const BOUNDARY_CONDITIONS<TV>& bc,const ARRAY<T,FACE_INDEX<TV::m>>* beta)
{
    typedef VECTOR<int,TV::m> IV;
    RANGE<IV> solve_domain=bc.grid.Cell_Indices();
    
    lm.filled_region_colors.Resize(solve_domain.Thickened(1),init_all,FLOOD_FILL::uncolorable);
    for(RANGE_ITERATOR<TV::m> it(solve_domain,0,0);it.Valid();it.Next())
    {
        if(bc.psi_D && (*bc.psi_D)(it.index)) lm.filled_region_colors(it.index)=FLOOD_FILL::uncolorable;
        else lm.filled_region_colors(it.index)=FLOOD_FILL::colorable;
    }
    ARRAY<int> filled_region_touches_dirichlet;
    // lm.filled_region_touches_dirichlet.Remove_All();
    
    lm.number_of_regions=Flood_Fill(solve_domain,at_cells,lm.filled_region_colors,bc.psi_N,&filled_region_touches_dirichlet);

    lm.matrix_index_to_cell_index_array.Resize(lm.number_of_regions);
    lm.cell_index_to_matrix_index.Resize(solve_domain.Thickened(1),init_new,-1);
    lm.A_array.Resize(lm.number_of_regions);
    lm.enforce_compatibility.Resize(lm.number_of_regions,init_new,true);
    rhs.Resize(lm.number_of_regions);
   
    //Compute_Matrix_Indices
    for(RANGE_ITERATOR<TV::m> it(solve_domain,1,0);it.Valid();it.Next())
    {
        int color=lm.filled_region_colors(it.index);
        if(color>=0) lm.cell_index_to_matrix_index(it.index)=lm.matrix_index_to_cell_index_array(color).Append(it.index);
    }
    for(int side=0;side<2*TV::m;side++) // Periodic cell indices
    {
        if(bc.bc(side)!=BOUNDARY_CONDITIONS<TV>::BC_PERIODIC) continue;
        for(RANGE_ITERATOR<TV::m> it(solve_domain,1,0,RI::ghost,side);it.Valid();it.Next())
            lm.cell_index_to_matrix_index(it.index)=lm.cell_index_to_matrix_index(wrap(it.index,bc.grid.Cell_Indices().min_corner,bc.grid.Cell_Indices().max_corner));
    }
    //End Compute_Matrix

    // Find A
    auto Get_Beta=[beta](FACE_INDEX<TV::m> face){ if(beta) return (*beta)(face); else return (T)1;};
    
    for(int i=0;i<lm.A_array.m;i++)
    {
        lm.A_array(i).Reset(lm.matrix_index_to_cell_index_array(i).m);
        rhs(i).Resize(lm.matrix_index_to_cell_index_array(i).m);
    }

    TV one_over_dx2=Inverse(bc.grid.dX*bc.grid.dX);
    for(RANGE_ITERATOR<TV::m> it(solve_domain,1,0);it.Valid();it.Next())
    {
        int color=lm.filled_region_colors(it.index);
        assert(color!=-1); // everything has been colored
        if(color<0) continue;
        
        int matrix_index=lm.cell_index_to_matrix_index(it.index);
        SPARSE_MATRIX_FLAT_MXN<T>& A=lm.A_array(color);
        ARRAY<T>& b=rhs(color);
        b(matrix_index)=f(it.index);

        T diagonal=0;
        for(int axis=0;axis<TV::m;axis++)
        {
            for(int s=0;s<2;s++)
            {
                FACE_INDEX<TV::m> neighbor_face(axis,it.index.Add_Axis(axis,s));
                IV neighbor_cell=it.index.Add_Axis(axis,2*s-1);
                if(lm.filled_region_colors.Valid_Index(neighbor_cell)) 
                {
                    if(bc.psi_N && (*bc.psi_N)(neighbor_face))
                    {
                        if(bc.psi_N_values) b(matrix_index)+=Get_Beta(neighbor_face)*(*bc.psi_N_values)(neighbor_face)*bc.grid.one_over_dX(axis);
                    }
                    else
                    {
                        T element=Get_Beta(neighbor_face)*one_over_dx2[axis];
                        diagonal+=element;

                        int neighbor_index=lm.cell_index_to_matrix_index(neighbor_cell);
                        if(neighbor_index<0)
                        {
                            if(bc.psi_D_values)
                                b(matrix_index)-=element*(*bc.psi_D_values)(neighbor_cell);
                            lm.enforce_compatibility(color)=false;
                        }
                        else
                            A.Append_Entry_To_Current_Row(neighbor_index,element);
                    }
                }
            }
        }
        A.Append_Entry_To_Current_Row(matrix_index,-diagonal);
        A.Finish_Row();
    }
    for(int i=0;i<lm.A_array.m;i++)
        lm.A_array(i).Sort_Entries();
}
//#####################################################################
// Function Compute_Laplace_Matrix
//#####################################################################
template<class TV,class T,class IV>
void Compute_Laplace_Matrix(SPARSE_MATRIX_FLAT_MXN<T>& A,ARRAY<T>& rhs,
    const ARRAYS_ND_BASE<T,VECTOR<int,TV::m>>& f,const BOUNDARY_CONDITIONS<TV>& bc,
    ARRAY<IV>& matrix_index_to_cell_index_array,ARRAY<int,IV>& cell_index_to_matrix_index,
    bool& enforce_compatibility,const ARRAY<T,FACE_INDEX<TV::m>>* beta)
{
    RANGE<IV> solve_domain=bc.grid.Cell_Indices();
    cell_index_to_matrix_index.Resize(solve_domain.Thickened(1),init_new,-1);
    enforce_compatibility=true;

    for(RANGE_ITERATOR<TV::m> it(solve_domain,0,0);it.Valid();it.Next())
    {
        if(!bc.psi_D || !(*bc.psi_D)(it.index))
            cell_index_to_matrix_index(it.index)=matrix_index_to_cell_index_array.Append(it.index);
    }
    for(int side=0;side<2*TV::m;side++) // Periodic cell indices
    {
        if(bc.bc(side)!=BOUNDARY_CONDITIONS<TV>::BC_PERIODIC) continue;
        for(RANGE_ITERATOR<TV::m> it(solve_domain,1,0,RI::ghost,side);it.Valid();it.Next())
            cell_index_to_matrix_index(it.index)=cell_index_to_matrix_index(wrap(it.index,bc.grid.Cell_Indices().min_corner,bc.grid.Cell_Indices().max_corner));
    }
    //End Compute_Matrix

    // Find A
    auto Get_Beta=[beta](FACE_INDEX<TV::m> face){ if(beta) return (*beta)(face); else return (T)1;};
    
    A.Reset(matrix_index_to_cell_index_array.m);
    rhs.Resize(matrix_index_to_cell_index_array.m);

    TV one_over_dx2=Inverse(bc.grid.dX*bc.grid.dX);
    for(RANGE_ITERATOR<TV::m> it(solve_domain,1,0);it.Valid();it.Next())
    {
        int matrix_index=cell_index_to_matrix_index(it.index);
        if(matrix_index<0) continue;
        
        rhs(matrix_index)=f(it.index);

        T diagonal=0;
        for(int axis=0;axis<TV::m;axis++)
        {
            for(int s=0;s<2;s++)
            {
                FACE_INDEX<TV::m> neighbor_face(axis,it.index.Add_Axis(axis,s));
                IV neighbor_cell=it.index.Add_Axis(axis,2*s-1);
                if(cell_index_to_matrix_index.Valid_Index(neighbor_cell)) 
                {
                    if(bc.psi_N && (*bc.psi_N)(neighbor_face))
                    {
                        if(bc.psi_N_values)
                            rhs(matrix_index)+=Get_Beta(neighbor_face)*(*bc.psi_N_values)(neighbor_face)*bc.grid.one_over_dX(axis);
                    }
                    else
                    {
                        T element=Get_Beta(neighbor_face)*one_over_dx2[axis];
                        diagonal+=element;

                        int neighbor_index=cell_index_to_matrix_index(neighbor_cell);
                        if(neighbor_index<0)
                        {
                            if(bc.psi_D_values)
                                rhs(matrix_index)-=element*(*bc.psi_D_values)(neighbor_cell);
                            enforce_compatibility=false;
                        }
                        else
                            A.Append_Entry_To_Current_Row(neighbor_index,element);
                    }
                }
            }
        }
        A.Append_Entry_To_Current_Row(matrix_index,-diagonal);
        A.Finish_Row();
    }
    A.Sort_Entries();
}

// template struct LAPLACE_MATRIX<VECTOR<float,2>>;
// template struct LAPLACE_MATRIX<VECTOR<float,3>>;
// template struct LAPLACE_MATRIX<VECTOR<double,2>>;
// template struct LAPLACE_MATRIX<VECTOR<double,3>>;

// template struct BOUNDARY_CONDITIONS<VECTOR<float,2>>;
// template struct BOUNDARY_CONDITIONS<VECTOR<float,3>>;
// template struct BOUNDARY_CONDITIONS<VECTOR<double,2>>;
// template struct BOUNDARY_CONDITIONS<VECTOR<double,3>>;

template void Compute_Laplace_Matrix<VECTOR<double, 2>, double>(LAPLACE_MATRIX<VECTOR<double, 2> >&, ARRAY<ARRAY<double, int>, int>&, ARRAYS_ND_BASE<double, VECTOR<int, VECTOR<double, 2>::m> > const&, BOUNDARY_CONDITIONS<VECTOR<double, 2> > const&, ARRAY<double, FACE_INDEX<VECTOR<double, 2>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<double, 3>, double>(LAPLACE_MATRIX<VECTOR<double, 3> >&, ARRAY<ARRAY<double, int>, int>&, ARRAYS_ND_BASE<double, VECTOR<int, VECTOR<double, 3>::m> > const&, BOUNDARY_CONDITIONS<VECTOR<double, 3> > const&, ARRAY<double, FACE_INDEX<VECTOR<double, 3>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<float, 2>, float>(LAPLACE_MATRIX<VECTOR<float, 2> >&, ARRAY<ARRAY<float, int>, int>&, ARRAYS_ND_BASE<float, VECTOR<int, VECTOR<float, 2>::m> > const&, BOUNDARY_CONDITIONS<VECTOR<float, 2> > const&, ARRAY<float, FACE_INDEX<VECTOR<float, 2>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<float, 3>, float>(LAPLACE_MATRIX<VECTOR<float, 3> >&, ARRAY<ARRAY<float, int>, int>&, ARRAYS_ND_BASE<float, VECTOR<int, VECTOR<float, 3>::m> > const&, BOUNDARY_CONDITIONS<VECTOR<float, 3> > const&, ARRAY<float, FACE_INDEX<VECTOR<float, 3>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<double,2>,double,VECTOR<int,2> >(SPARSE_MATRIX_FLAT_MXN<double>&,ARRAY<double,int>&,ARRAYS_ND_BASE<double,VECTOR<int,VECTOR<double,2>::m> > const&,BOUNDARY_CONDITIONS<VECTOR<double,2> > const&,ARRAY<VECTOR<int,2>,int>&,ARRAY<int,VECTOR<int,2> >&,bool&,ARRAY<double,FACE_INDEX<VECTOR<double,2>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<float,2>,float,VECTOR<int,2> >(SPARSE_MATRIX_FLAT_MXN<float>&,ARRAY<float,int>&,ARRAYS_ND_BASE<float,VECTOR<int,VECTOR<float,2>::m> > const&,BOUNDARY_CONDITIONS<VECTOR<float,2> > const&,ARRAY<VECTOR<int,2>,int>&,ARRAY<int,VECTOR<int,2> >&,bool&,ARRAY<float,FACE_INDEX<VECTOR<float,2>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<double,3>,double,VECTOR<int,3> >(SPARSE_MATRIX_FLAT_MXN<double>&,ARRAY<double,int>&,ARRAYS_ND_BASE<double,VECTOR<int,VECTOR<double,3>::m> > const&,BOUNDARY_CONDITIONS<VECTOR<double,3> > const&,ARRAY<VECTOR<int,3>,int>&,ARRAY<int,VECTOR<int,3> >&,bool&,ARRAY<double,FACE_INDEX<VECTOR<double,3>::m> > const*);
template void Compute_Laplace_Matrix<VECTOR<float,3>,float,VECTOR<int,3> >(SPARSE_MATRIX_FLAT_MXN<float>&,ARRAY<float,int>&,ARRAYS_ND_BASE<float,VECTOR<int,VECTOR<float,3>::m> > const&,BOUNDARY_CONDITIONS<VECTOR<float,3> > const&,ARRAY<VECTOR<int,3>,int>&,ARRAY<int,VECTOR<int,3> >&,bool&,ARRAY<float,FACE_INDEX<VECTOR<float,3>::m> > const*);
}
