#include <Core/Poisson/HELMHOLTZ.h>
#include <Core/Arrays/FACE_ARRAYS.h>
#include <Core/Arrays/FLOOD_FILL.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/Krylov_Solvers/CONJUGATE_GRADIENT.h>
#include <Core/Krylov_Solvers/KRYLOV_SOLVER.h>
#include <Core/Krylov_Solvers/KRYLOV_VECTOR_BASE.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/Poisson/COMPUTE_LAPLACE_MATRIX.h>
#include <Core/Poisson/LAPLACE.h>
#include <Core/Poisson/POISSON_UNIFORM.h>
#include <Core/Read_Write/OCTAVE_OUTPUT.h>
using namespace PhysBAM;
//#####################################################################
// Constructor
//#####################################################################
template<class TV> HELMHOLTZ<TV>::
HELMHOLTZ(const GRID<TV>& grid_input,bool n)
    :grid(grid_input),nodes(n)
{
}
//#####################################################################
// Function Solve
//#####################################################################
template<class TV> bool HELMHOLTZ<TV>::
Solve(ARRAYS_ND_BASE<T,IV>& u,const ARRAYS_ND_BASE<T,IV>& f,RANGE<IV> grid_domain,PCG_OPTIONS<T>& options)
{
    BOUNDARY_CONDITIONS<TV> bc;
    bc.psi_N=psi_N;
    bc.psi_N_values=psi_N_values;
    bc.psi_D=psi_D;
    bc.psi_D_values=psi_D_values;
    bc.grid=grid;
    for(int side=0;side<2*TV::m;side++)
    {
        if(periodic_boundary(side/2)) bc.bc(side)=BOUNDARY_CONDITIONS<TV>::BC_PERIODIC;
        else bc.bc(side)=BOUNDARY_CONDITIONS<TV>::BC_NEUMANN;
    }

    LAPLACE_MATRIX<TV> lm;
    
    RANGE<IV> valid_domain(grid_domain);
    // if(nodes) for(int a=0;a<TV::m;a++)
    //         if(periodic_boundary(a)) valid_domain.max_corner(a)--;
    RANGE<IV> solve_domain=valid_domain.Intersect(u.domain);
    ARRAY<ARRAY<T>> b_array;
    auto* invert_beta=new ARRAY<T,FACE_INDEX<TV::m> >(grid_domain);
    for(FACE_RANGE_ITERATOR<TV::m> it(grid_domain);it.Valid();it.Next())
    {
        if(beta) (*invert_beta)(it.face)=-(*beta)(it.face);
        else (*invert_beta)(it.face)=-beta_c;
    }
    Compute_Laplace_Matrix<TV,T>(lm,b_array,f,bc,invert_beta);
    delete invert_beta;
    
    Find_A(lm,solve_domain);
    
    bool converged=true;
    for(int color=0;color<lm.number_of_regions;color++)
        if(lm.matrix_index_to_cell_index_array(color).m>0)
        {
            options.enforce_compatibility=lm.enforce_compatibility(color);
            converged=Solve_Subregion(lm,lm.A_array(color),b_array(color),u,options,color);
        }

    if(psi_D && psi_D_values)
        for(RANGE_ITERATOR<TV::m> it(solve_domain,0,0);it.Valid();it.Next())
            if((*psi_D)(it.index)) u(it.index)=(*psi_D_values)(it.index);

    // if(nodes) Fill_Ghost_Periodic<T,IV>(valid_domain.Edge_Lengths(),u,periodic_boundary);
    return converged;
}
//#####################################################################
// Function Find_A
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Find_A(LAPLACE_MATRIX<TV>& lm,const RANGE<IV>& u_size)
{
    auto Get_Alpha=[this](IV index){ if(alpha) return (*alpha)(index); else return alpha_c;};    

    for(int i=0;i<lm.A_array.m;i++)
        lm.A_array(i).Initialize_Diagonal_Index();
    
    // TV one_over_dx2=Inverse(grid.dX*grid.dX);
    for(RANGE_ITERATOR<TV::m> it(u_size,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);

        T diagonal=Get_Alpha(it.index);
        if(diagonal) lm.enforce_compatibility(color)=false;
        A.A(A.diagonal_index(matrix_index)).a+=diagonal;
    }
}
//#####################################################################
// Function Solve_Subregion
//#####################################################################
template<class TV> bool HELMHOLTZ<TV>::
Solve_Subregion(LAPLACE_MATRIX<TV>& lm,SPARSE_MATRIX_FLAT_MXN<T>& A,ARRAY<T>& b,ARRAYS_ND_BASE<T,IV>& u,const PCG_OPTIONS<T>& options,const int color)
{
    MATRIX_SYSTEM<SPARSE_MATRIX_FLAT_MXN<T>,T,KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > > system(A);

    // Set up nullspace
    ARRAY<T> nullspace(A.m,init_new,1);
    ARRAY<T> diag_sqrt_inv;
    if(diagonal_scaling)
    {
        LOG::cout<<"scaling diagonals before preconditioner"<<std::endl;
        A.Initialize_Diagonal_Index();
        diag_sqrt_inv.Resize(A.diagonal_index.m);
        for(int i=0;i<A.diagonal_index.m;i++)
        {
            T square_root=sqrt(A.A(A.diagonal_index(i)).a);
            diag_sqrt_inv(i)=1/square_root;
        }
        A.Set_Diagonal_Times(diag_sqrt_inv);
        A.Set_Times_Diagonal(diag_sqrt_inv);
        b*=diag_sqrt_inv;
        nullspace=diag_sqrt_inv*nullspace*diag_sqrt_inv;
    }    
    nullspace.Normalize(); // nullspace "must be orthogonal and normalized"
    KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > null(nullspace);
    if(options.enforce_compatibility) system.nullspace_vectors.Append(&null);

    // Set up preconditioner
    KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > temp(nullspace);
    bool incomplete_cholesky=true;
    if(incomplete_cholesky)
    {
        system.use_preconditioner=true;
        system.preconditioner_commutes_with_projection=false;
        bool recompute_preconditioner=true;
        if(recompute_preconditioner||!A.C)
        {
            T modified_incomplete_cholesky_coefficient=(T).97;
            T preconditioner_zero_tolerance=(T)1e-8;
            T preconditioner_zero_replacement=(T)1e-8;
            bool modified_incomplete_cholesky=true;
            A.Construct_Incomplete_Cholesky_Factorization(modified_incomplete_cholesky,
                modified_incomplete_cholesky_coefficient,preconditioner_zero_tolerance,
                preconditioner_zero_replacement);
        }
        system.Set_Preconditioner(*A.C,temp);
    }

    // Set up solver
    CONJUGATE_GRADIENT<T> solver;
    solver.print_diagnostics=options.show_results;
    solver.print_residuals=options.show_residual;
    solver.restart_iterations=options.restart_iterations;
    solver.ensure_convergence=options.ensure_convergence;
    if(print_system) solver.print_matrix=octave_output.c_str();

    int number_of_unknowns=lm.matrix_index_to_cell_index_array(color).m;
    T tolerance=max(options.relative_tolerance*b.Max_Abs(),options.absolute_tolerance);
    if(options.show_results) LOG::cout<<"solving "<<number_of_unknowns<<" cells to tolerance "<<tolerance<<std::endl;

    ARRAY<T> x(number_of_unknowns);
    for(int i=0;i<number_of_unknowns;i++) x(i)=u(lm.matrix_index_to_cell_index_array(color)(i));
    ARRAY<KRYLOV_VECTOR_BASE<T>*> vectors;
    KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > kx;
    KRYLOV_VECTOR_WRAPPER<T,ARRAY<T> > kb;
    std::swap(kx.v,x);
    std::swap(kb.v,b);

    int desired_iterations=options.max_iterations?options.max_iterations:(A.n-options.enforce_compatibility);
    bool converged=solver.Solve(system,kx,kb,vectors,tolerance,0,desired_iterations);
    if(options.enforce_compatibility) system.Project_Nullspace(kx);
    
    std::swap(kx.v,x);
    std::swap(kb.v,b);

    if(diagonal_scaling) x*=diag_sqrt_inv;

    if(print_system)
    {
        static int id=0;
        auto name=LOG::sprintf("%s-%i.txt",octave_output.c_str(),id++);
        OCTAVE_OUTPUT<T> o_out(name,true);
        o_out.Write("x",x);
        o_out.Write("p",nullspace);
    }
    
    for(int i=0;i<number_of_unknowns;i++) // move solution to u
        u(lm.matrix_index_to_cell_index_array(color)(i))=x(i);
    
    vectors.Delete_Pointers_And_Clean_Memory();
    return converged;
}
//#####################################################################
// Function Set_Neumann_Boundaries
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Set_Neumann_Boundaries(const ARRAY<bool,FACE_INDEX<TV::m> >& N, const ARRAY<T,FACE_INDEX<TV::m> >* N_values)
{
    psi_N=&N;
    psi_N_values=N_values;
}
//#####################################################################
// Function Set_Dirichlet_Boundaries
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Set_Dirichlet_Boundaries(const ARRAYS_ND_BASE<bool,IV>& D, const ARRAYS_ND_BASE<T,IV>* D_values)
{
    psi_D=&D;
    psi_D_values=D_values;
}
//#####################################################################
// Function Set_Periodic_Boundaries
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Set_Periodic_Boundaries(const VECTOR<bool,TV::m>& p)
{
    periodic_boundary=p;
}
//#####################################################################
// Function Set_Beta
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Set_Beta(T b)
{
    beta_c=b;
    beta=0;
}
template<class TV> void HELMHOLTZ<TV>::
Set_Beta(const ARRAY<T,FACE_INDEX<TV::m> >& b)
{
    beta=&b;
}
//#####################################################################
// Function Set_Alpha
//#####################################################################
template<class TV> void HELMHOLTZ<TV>::
Set_Alpha(T a)
{
    alpha_c=a;
    alpha=0;
}
template<class TV> void HELMHOLTZ<TV>::
Set_Alpha(const ARRAYS_ND_BASE<T,IV>& a)
{
    alpha=&a;
}
//#####################################################################
namespace PhysBAM{
//template class HELMHOLTZ<VECTOR<float,1> >;
template class HELMHOLTZ<VECTOR<float,2> >;
template class HELMHOLTZ<VECTOR<float,3> >;
//template class HELMHOLTZ<VECTOR<double,1> >;
template class HELMHOLTZ<VECTOR<double,2> >;
template class HELMHOLTZ<VECTOR<double,3> >;
template struct PCG_OPTIONS<float>;
template struct PCG_OPTIONS<double>;
}
