//#####################################################################
// Copyright 2009-2010, Michael Lentine, Andrew Selle.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Geometry_Particles/DEBUG_PARTICLES.h>
#include <Core/Log/DEBUG_SUBSTEPS.h>
#include <Core/Log/LOG.h>
#include <Core/Log/SCOPE.h>
#include "SMOKE_DRIVER.h"
#include "SMOKE_EXAMPLE.h"

using namespace PhysBAM;
//#####################################################################
// Constructor
//#####################################################################
template<class TV> SMOKE_DRIVER<TV>::
SMOKE_DRIVER(SMOKE_EXAMPLE<TV>& example)
    :ghost(5),example(example)
{
    DEBUG_SUBSTEPS::write_substeps_level=example.write_substeps_level;
    DEBUG_SUBSTEPS::writer=[this](const std::string& title,int level){this->example.viewer_output.Write_Substep(title,level);};
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV> SMOKE_DRIVER<TV>::
~SMOKE_DRIVER()
{
    DEBUG_SUBSTEPS::writer=0;
}
//#####################################################################
// Function Execute_Main_Program
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Execute_Main_Program()
{
    Initialize();
    Simulate_To_Frame(example.last_frame);
}
//#####################################################################
// Function Initialize
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Initialize()
{
    DEBUG_SUBSTEPS::write_substeps_level=example.write_substeps_level;

    // setup time
    if(example.restart) current_frame=example.restart;else current_frame=0;    
    time=example.Time_At_Frame(current_frame);

    example.boundary=&example.boundary_scalar;

    // setup grids and velocities
    example.projection.Initialize_Grid(example.grid);
    example.face_velocities.Resize(example.grid.Cell_Indices());
    example.density.Resize(example.grid.Cell_Indices(ghost));
    example.temperature.Resize(example.grid.Cell_Indices(ghost));  // add temperature
    example.Initialize_Fields();

    // setup laplace
    example.projection.elliptic_solver->Set_Relative_Tolerance(1e-9);
    example.projection.elliptic_solver->pcg.Set_Maximum_Iterations(1000);
    example.projection.elliptic_solver->pcg.evolution_solver_type=krylov_solver_cg;
    example.projection.elliptic_solver->pcg.cg_restart_iterations=40;

    if(example.restart) example.Read_Output_Files();
    
    // setup domain boundaries
    VECTOR<VECTOR<bool,2>,TV::m> constant_extrapolation;constant_extrapolation.Fill(VECTOR<bool,2>::Constant_Vector(true));
    example.boundary->Set_Constant_Extrapolation(constant_extrapolation);
    example.Set_Boundary_Conditions(time, 1); // get so CFL is correct

    example.Register_Output_Files();
    if(!example.restart) example.viewer_output.Flush_Frame("");
    output_number=0;
}

//#####################################################################
// Function Add_Buoyancy_Force
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Add_Buoyancy_Force(const T dt,const T time)
{
    for(FACE_ITERATOR<TV> iterator(example.grid);iterator.Valid();iterator.Next()){
        if(iterator.face.axis==1){
            IV c1,c2;
            example.grid.Cells_Touching_Face(iterator.face.axis,iterator.Face_Index(),c1,c2);
            T rho=(example.density(c1)+example.density(c2))*(T).5;
            T tem=(example.temperature(c1)+example.temperature(c2))*(T).5; // add temperature 
            example.face_velocities(iterator.Full_Index())+=-dt*example.alpha*rho+dt*example.beta*tem;}}// add temperature
}
//#####################################################################
// Function Scalar_Advance
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Scalar_Advance(const T dt,const T time)
{  
    example.Get_Scalar_Field_Sources(time);
    ARRAY<T,IV> density_ghost(example.grid.Cell_Indices(ghost));
    ARRAY<T,IV> temperature_ghost(example.grid.Cell_Indices(ghost));  // add temperature
    example.boundary->Set_Fixed_Boundary(true,0);
    example.boundary->Fill_Ghost_Cells(example.grid,example.density,density_ghost,ghost);
    example.boundary->Fill_Ghost_Cells(example.grid,example.temperature,temperature_ghost,ghost);
    example.advection_scalar.Update_Advection_Equation_Cell(example.grid,example.density,density_ghost,example.face_velocities,dt,time);
    example.advection_scalar.Update_Advection_Equation_Cell(example.grid,example.temperature,temperature_ghost,example.face_velocities,dt,time);
    example.boundary->Set_Fixed_Boundary(false);
}
//#####################################################################
// Function Convect
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Convect(const T dt,const T time)
{
    example.boundary->Set_Fixed_Boundary(true,0);
    ARRAY<T,FACE_INDEX<TV::m> > face_velocities_ghost(example.grid.Cell_Indices(ghost),no_init);
    example.boundary->Fill_Ghost_Faces(example.grid,example.face_velocities,face_velocities_ghost,ghost);
    example.advection_scalar.Update_Advection_Equation_Face(
        example.grid,example.face_velocities,face_velocities_ghost,face_velocities_ghost,dt,time);
    example.boundary->Set_Fixed_Boundary(false);
}
//#####################################################################
// Function Project
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Project(const T dt,const T time)
{
    example.Set_Boundary_Conditions(time+dt,1);  // set 0 then it doesn't have velocity after initial
    example.projection.p.array*=dt; // rescale pressure for guess
    example.boundary->Apply_Boundary_Condition_Face(example.grid,example.face_velocities);
    example.projection.Make_Divergence_Free(example.face_velocities,dt,time);
    example.projection.p.array*=(1/dt); // unscale pressure
    LOG::cout << "projection" << std::endl;
}

//#####################################################################
// Function Advance_To_Target_Time
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Advance_To_Target_Time(const T target_time)
{
    bool done=false;for(int substep=1;!done;substep++){
        // choose time step
        LOG::Time("Substep");
        T dt=example.cfl*example.CFL(example.face_velocities);        
        if(time+dt>=target_time){dt=target_time-time;done=true;}
        else if(time+2*dt>=target_time){dt=.5*(target_time-time);}
        Scalar_Advance(dt,time);
        Convect(dt,time);
        Add_Buoyancy_Force(dt,time);
        Print_Max_Divergence("Before projection");
        Project(dt,time);
        Print_Max_Divergence("After projection");
        time+=dt;}
}
//#####################################################################
// Function Simulate_To_Frame
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Simulate_To_Frame(const int frame)
{
    while(current_frame<frame){
        LOG::SCOPE scope("FRAME","Frame %d",current_frame+1);
        Advance_To_Target_Time(example.Time_At_Frame(current_frame+1));
        example.viewer_output.Flush_Frame("");
        current_frame++;}
}
//#####################################################################
// Function Print_Max_Divergence
//#####################################################################
template<class TV> void SMOKE_DRIVER<TV>::
Print_Max_Divergence(const char* str)
{
    if(!example.debug_divergence) return;
    T max_div=(T)0;
    for(CELL_ITERATOR<TV> iterator(example.grid);iterator.Valid();iterator.Next()){
        T d=0;
        for(int axis=0;axis<TV::m;axis++)
            d+=example.face_velocities(FACE_INDEX<TV::m>(axis,iterator.Second_Face_Index(axis)))
                -example.face_velocities(FACE_INDEX<TV::m>(axis,iterator.First_Face_Index(axis)));
        T ad=abs(d); if(ad>max_div) max_div=ad;}
    LOG::cout<<str<<" max(div(v)) "<<max_div<<std::endl;
}
//#####################################################################



namespace PhysBAM{
template class SMOKE_DRIVER<VECTOR<float,2> >;
template class SMOKE_DRIVER<VECTOR<float,3> >;
template class SMOKE_DRIVER<VECTOR<double,2> >;
template class SMOKE_DRIVER<VECTOR<double,3> >;
}
