//#####################################################################
// 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 <Core/Projection/CUT_CELL_PROJECTION.h>
#include <Core/Basic_Geometry/BASIC_SIMPLEX_POLICY.h>
#include <Core/Basic_Geometry/SEGMENT_2D.h>
#include <Core/Basic_Geometry/TRIANGLE_3D.h>
#include <Core/Boundaries/BOUNDARY_PERIODIC.h>
#include <Core/Boundaries/PERIODIC_BOUNDARY_HELPER.h>
#include <Core/Geometry_Particles/DEBUG_PARTICLES.h>
#include <Core/Grids/FACE_RANGE_ITERATOR.h>
#include <Core/Grids_Uniform_Computations/MARCHING_CUBES.h>
#include <Core/Krylov_Solvers/CONJUGATE_GRADIENT.h>
#include <Core/Krylov_Solvers/KRYLOV_VECTOR_WRAPPER.h>
#include <Core/Krylov_Solvers/MATRIX_SYSTEM.h>
#include <Core/Log/LOG.h>
#include <Core/Matrices/SPARSE_MATRIX_FLAT_MXN.h>
#include <Core/Read_Write/OCTAVE_OUTPUT.h>
namespace PhysBAM
{
//#####################################################################
// Function Cut_Cell_Projection
//#####################################################################
template<class TV> void CUT_CELL_PROJECTION<TV>::
Cut_Cell_Projection(const GRID<TV>& grid,
    ARRAY<T,FACE_INDEX<TV::m> >& u,T density,T dt)
{
    typedef typename BASIC_SIMPLEX_POLICY<TV,TV::m>::SIMPLEX_FACE FACE;
    typedef KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > VEC;
    typedef MATRIX_SYSTEM<SPARSE_MATRIX_FLAT_MXN<T>,T,VEC> MAT;

    // TODO: how much mass does the surface velocity edge get?  mass -> 0?
    
    // TODO: fill in valid xfer data flags.

    static const int num_corners=1<<TV::m;
    static const int partial_object_flag=1<<29;
    static const int partial_surface_flag=1<<30;
    static const int index_mask=partial_object_flag-1;

    ARRAY<int,IV> cell_index(u.domain.Thickened(1),init_new,-9);
    VEC rhs,sol,null,precon_tmp;
    SPARSE_MATRIX_FLAT_MXN<T> A;
    MAT sys(A);

    PHYSBAM_ASSERT(!surface_phi || surface_phi->domain.Contains(u.domain));
    PHYSBAM_ASSERT(!object_phi || object_phi->domain.Contains(u.domain));
    
    // nonnegative = valid
    // -1 = air
    // -2 = object
    // -9 = uninitialized
    for(int s=0;s<2*TV::m;s++){
        if(bc_type(s)==0) continue;
        int value=-bc_type(s);
        for(RANGE_ITERATOR<TV::m> it(cell_index.domain,grid.Cell_Indices(),RI::ghost,s);it.Valid();it.Next())
            cell_index(it.index)=value;}

    ARRAY<IV> dof_cells;
    for(RANGE_ITERATOR<TV::m> it(u.domain);it.Valid();it.Next()){
        int& ci=cell_index(it.index);
        if(ci!=-9) continue;
        VECTOR<T,num_corners> phis;
        bool has_surface=false,has_object=false;
        if(object_phi){
            MARCHING_CUBES<TV>::Compute_Phis_For_Cell(phis,*object_phi,it.index);
            if(phis.Min()>=0){ci=-2;continue;} // all object
            has_object=phis.Max()>0;}
        if(surface_phi){
            MARCHING_CUBES<TV>::Compute_Phis_For_Cell(phis,*surface_phi,it.index);
            if(phis.Min()>=0){ci=-1;continue;} // all air
            has_surface=phis.Max()>0;}
        PHYSBAM_ASSERT(!has_object || !has_surface);

        ci=dof_cells.Append(it.index);


        if(has_surface) ci|=partial_surface_flag;
        if(has_object) ci|=partial_object_flag;

    }
    rhs.v.Resize(dof_cells.m,init_all,0);

    VECTOR<bool,TV::m> is_periodic;
    for(int i=0;i<TV::m;i++) is_periodic(i)=!bc_type(2*i);
    if(is_periodic.Contains(true))
        Fill_Ghost_Cells_Periodic(grid,cell_index,is_periodic);
    
    ARRAY<VECTOR<TV,TV::m> > surface;
    VECTOR<ARRAY<VECTOR<TV,TV::m> >,2*TV::m> boundary;
    VECTOR<VECTOR<ARRAY<VECTOR<TV,TV::m> >*,2>,2*TV::m> pboundary;
    for(int s=0;s<2*TV::m;s++) pboundary(s)(1)=&boundary(s);
    
    // Assumes square cells
    T one_over_dx=grid.one_over_dX(0);
    TV base_entry=grid.one_over_dX*grid.one_over_dX/density;
    // times 2 to compensate for pressure at face
    T base_d_entry=one_over_dx*one_over_dx/density*2;
    TV sc=grid.one_over_dX/density;
    bool has_nullspace=true;

    if(face_fractions)
        face_fractions->Fill(0);
    A.Reset(dof_cells.m);
    for(int i=0;i<dof_cells.m;i++)
    {
        IV index=dof_cells(i);
        int this_index=cell_index(index);
        bool has_surface=this_index&partial_surface_flag;
        bool has_object=this_index&partial_object_flag;
        this_index&=index_mask;
        
        VECTOR<T,TV::m*2> face_fraction;
        TV surface_area_weighted_normal;
        TV centroid;
        T total_area=0;
        T diag_entry=0;
        if(has_surface || has_object){
            VECTOR<T,num_corners> phis;
            if(has_surface)
                MARCHING_CUBES<TV>::Compute_Phis_For_Cell(phis,*surface_phi,index);
            if(has_object)
                MARCHING_CUBES<TV>::Compute_Phis_For_Cell(phis,*object_phi,index);
            surface.Remove_All();
            for(int s=0;s<2*TV::m;s++) boundary(s).Remove_All();
            MARCHING_CUBES<TV>::Get_Elements_For_Cell(surface,pboundary,phis);

            for(int s=0;s<2*TV::m;s++)
                for(const auto& t:boundary(s)){
                    face_fraction(s) += FACE::Size(t);
                }
            T flux=0;
            for(const auto& t:surface)
            {
                auto wt(t);
                for(auto&x:wt) x=grid.Cell_Domain(index).Point_From_Normalized_Coordinates(x);
                TV an=FACE::Area_Weighted_Normal(t);
                T area=an.Magnitude();
                surface_area_weighted_normal+=an;
                centroid += t.Sum()/t.m * area;
                total_area+=area;
                if(ave_v) flux+=an.Dot(ave_v(wt));
            }
            centroid/=total_area;
            centroid=grid.Cell_Domain(index).Point_From_Normalized_Coordinates(centroid);
            PHYSBAM_ASSERT(grid.Cell_Domain(index).Lazy_Inside(centroid));
            if(has_surface)
            {
                if(bc_p){
                    T p=bc_p(centroid)*dt;
                    rhs.v(this_index)+=base_d_entry*p*surface_area_weighted_normal.Magnitude();}
                // TODO: need rhs velocity at surface edge
                has_nullspace=false;
                diag_entry+=base_d_entry*surface_area_weighted_normal.Magnitude();
            }
            else
            {
                if(ave_v) rhs.v(this_index)-=one_over_dx*flux;
                else if(bc_v){
                    TV v=bc_v(centroid);
                    rhs.v(this_index)-=one_over_dx*surface_area_weighted_normal.Dot(v);
                }
            }
        }
        else face_fraction.Fill(1);

        if(face_fractions)
            for(int a=0;a<TV::m;a++)
                for(int s=0;s<2;s++)
                    (*face_fractions)({a, index.Add_Axis(a, s)}) = face_fraction(2 * a + s);

        for(int s=0;s<2*TV::m;s++)
        {
            if(!face_fraction(s)) continue;
            FACE_INDEX<TV::m> face(s/2,index);
            face.index(s/2)+=s%2;
            int sign=s%2?1:-1;
            IV cell(index);
            cell(s/2)+=sign;
            int ci=cell_index(cell);
            PHYSBAM_ASSERT(ci>=0 || ci==-2 || ci==-1 || !face_fraction(s));
            T entry=base_entry(s/2)*face_fraction(s);
            T div_entry=sign*grid.one_over_dX(s/2)*face_fraction(s);
            if(ci==-2)
            {
                T v=bc_v?bc_v(grid.Face_X(face))(s/2):0;
                if(!has_surface) rhs.v(this_index)-=div_entry*v;
                u(face)=v;
                continue;
            }
            if(ci==-1)
            {
                entry*=2; // compensate for pressure at face
                if(bc_p){
                    T p=bc_p(grid.Face_X(face))*dt;
                    rhs.v(this_index)+=entry*p;
                    u(face)-=sign*sc(s/2)*p;}
                has_nullspace=false;
            }
            else
            {
                A.Append_Entry_To_Current_Row(ci&index_mask,-entry);
            }
            if(!has_surface) rhs.v(this_index)-=div_entry*u(face);
            diag_entry+=entry;
        }
        A.Append_Entry_To_Current_Row(this_index,diag_entry);
        A.Finish_Row();
    }
    A.Sort_Entries();
    sol.v.Resize(rhs.v.m,init_all,0);
    if(has_nullspace){
        null.v.Resize(rhs.v.m,init_all,sqrt((T)1/rhs.v.m));
        sys.nullspace_vectors.Append(&null);}

    sys.use_preconditioner=use_preconditioner;
    sys.preconditioner_commutes_with_projection=(!has_nullspace||!use_preconditioner);
    if(use_preconditioner){
        A.Construct_Incomplete_Cholesky_Factorization();
        sys.P=A.C;
        precon_tmp.v.Resize(dof_cells.m,init_all,0);
        sys.temp_vector=&precon_tmp;}

    if(use_warm_start && pressure){
        if(valid_p) LOG::printf("%P\n",*valid_p);
        for(RANGE_ITERATOR<TV::m> it(valid_p->domain);it.Valid();it.Next()){
            int c=cell_index(it.index);
            if(c<0) continue;
            if(!valid_p || (*valid_p)(it.index))
                sol.v(c&index_mask)=(*pressure)(it.index)*dt;}}

    static int solve_id=-1;
    solve_id++;
    if(test_system) sys.Test_System(sol);
    OCTAVE_OUTPUT<T>* oo=0;
    if(print_matrix){
        oo=new OCTAVE_OUTPUT<T>(LOG::sprintf("M-%i.txt",solve_id).c_str());
        LOG::cout<<"solve id: "<<solve_id<<std::endl;
        oo->Write("M",sys,rhs);
        oo->Write("y",sol);
        oo->Write_Preconditioner("C",sys,rhs);
        oo->Write_Projection("P",sys,rhs);
        oo->Write("b",rhs);}

    ARRAY<KRYLOV_VECTOR_BASE<T>*> av;
    CONJUGATE_GRADIENT<T> cg;
    cg.finish_before_indefiniteness=true;
    cg.relative_tolerance=true;
    cg.print_diagnostics=true;
    cg.print_residuals=true;
    bool converged=cg.Solve(sys,sol,rhs,av,solver_tolerance,0,solver_iterations);
    if(!converged) LOG::printf("SOLVER DID NOT CONVERGE.\n");

    if(print_matrix){
        oo->Write("x",sol);
        delete oo;}

    if(print_residual){
        sys.Multiply(sol,*av(0));
        *av(0)-=rhs;
        T r=sys.Convergence_Norm(*av(0));
        LOG::cout<<"residual: "<<r<<std::endl;}

    for(FACE_RANGE_ITERATOR<TV::m> it(u.domain);it.Valid();it.Next()){
        int ca=cell_index(it.face.First_Cell_Index());
        int cb=cell_index(it.face.Second_Cell_Index());
        if(ca<0 && cb<0){
            if(valid_u) (*valid_u)(it.face)=false;
            continue;}
        if(valid_u) (*valid_u)(it.face)=true;
        if(ca==-2 || cb==-2) continue;
        T pa=ca<0?0:sol.v(ca&index_mask),pb=cb<0?0:sol.v(cb&index_mask);
        if(ca<0) pb*=2;
        if(cb<0) pa*=2;
        u(it.face)-=(pb-pa)*sc(it.face.axis);}

    if(valid_p || pressure){
        if(valid_p) valid_p->Resize(u.domain,no_init);
        if(pressure) pressure->Resize(u.domain,no_init);
        for(RANGE_ITERATOR<TV::m> it(valid_p->domain);it.Valid();it.Next()){
            int c=cell_index(it.index);
            if(c<0) continue;
            if(valid_p) (*valid_p)(it.index)=true;
            if(pressure) (*pressure)(it.index)=sol.v(c&index_mask)/dt;}}

    if(is_periodic.Contains(true)){
        Fill_Ghost_Faces_Periodic(grid,u,is_periodic);
        if(valid_u) Fill_Ghost_Faces_Periodic(grid,*valid_u,is_periodic);}
}
//#####################################################################
template struct CUT_CELL_PROJECTION<VECTOR<float,2> >;
template struct CUT_CELL_PROJECTION<VECTOR<float,3> >;
template struct CUT_CELL_PROJECTION<VECTOR<double,2> >;
template struct CUT_CELL_PROJECTION<VECTOR<double,3> >;
}
