#include <Core/Geometry_Particles/GEOMETRY_PARTICLES.h>
#include <Core/Grids/GRID.h>
#include <Core/Images/TEX_FILE.h>
#include <Core/Math_Tools/RANGE.h>
#include <Core/Math_Tools/RANGE_ITERATOR.h>
#include <Core/Parsing/PARSE_ARGS.h>
#include <Core/Particles/PARTICLES.h>
#include <Core/Read_Write/FILE_UTILITIES.h>
#include <Core/Symbolics/PROGRAM.h>
#include <Core/Symbolics/PROGRAM_CONTEXT.h>
#include <Core/Symbolics/PROGRAM_DEFINITION.h>
#include <Core/Utilities/VIEWER_DIR.h>
#include <fstream>
#include <iostream>

using namespace PhysBAM;

typedef double T;
typedef VECTOR<T,2> TV;
typedef VECTOR<int,2> IV;
typedef VECTOR<T,3> C;

struct COL_PROG
{
    std::string code;
    PROGRAM<T> prog;
    mutable PROGRAM_CONTEXT<T> context;
    C col;
    ARRAY<std::string,MEM_IN> var_in={"u","v","x","y"};
    ARRAY<std::string,MEM_OUT> var_out={"r","g","b"};

    void Init()
    {
        if(code.size())
        {
            Simple_Compile(prog,code,var_in,var_out);
            context.Initialize(prog);
        }
    }

    C Color(const TV& x,const TV& u) const
    {
        C c=col;
        if(code.size())
        {
            for(int i=0;i<TV::m;i++) context.data_in(MEM_IN(i))=u(i);
            for(int i=0;i<TV::m;i++) context.data_in(MEM_IN(i+TV::m))=x(i);
            prog.Execute(context);
            for(int i=0;i<3;i++) c(i)=context.data_out(MEM_OUT(i));
        }
        return c;
    }
};


int main(int argc,char *argv[])
{
    std::string sim="output";

    bool use_grid=false;
    RANGE<IV> bounds;
    bool custom_bounds=false;
    C grid_col{.5,.5,.5};
    T grid_width=0;

    bool use_u=false;
    COL_PROG col_u;
    T u_width=0;
    T u_scale=1;
    T u_min_len=0;
    bool u_arrow=false;
    std::string u_file="node_velocities";

    bool use_part=false;
    T part_size=.1;
    COL_PROG col_x;
    COL_PROG col_v;
    T part_v_width=0;
    T part_v_scale=1;
    T part_v_min_len=0;
    bool part_v_arrow=false;
    std::string part_file="mpm_particles";
    int first_frame=0;
    int last_frame=100000;
    int stride=1;
    
    T clamp=1e10;
    
    RANGE<TV> clip_domain;
    bool custom_clip=false;

    PARSE_ARGS parse_args(argc,argv);
    parse_args.Add("-i",&sim,"dir","simulation directory");

    parse_args.Add("-use_grid",&use_grid,"render grid");
    parse_args.Add("-grid_cells",&bounds,&custom_bounds,"domain","grid cell range to render");
    parse_args.Add("-grid_col",&grid_col,"color","grid color");
    parse_args.Add("-grid_width",&grid_width,"width","grid line width");

    parse_args.Add("-use_u",&use_u,"render node velocities");
    parse_args.Add("-u_col",&col_u.col,"color","fixed u color");
    parse_args.Add("-u_col_prog",&col_u.code,"color","program to compute u color (r=...;g=...;b=...;)");
    parse_args.Add("-u_width",&u_width,"width","u line width");
    parse_args.Add("-u_scale",&u_scale,"scale","scale u length");
    parse_args.Add("-u_min_len",&u_min_len,"min_len","min_len u");
    parse_args.Add("-u_arrow",&u_arrow,"draw velocity as arrow");
    parse_args.Add("-u_file",&u_file,"file","u filename");

    parse_args.Add("-use_part",&use_part,"render particles");
    parse_args.Add("-part_col",&col_x.col,"color","fixed particle color");
    parse_args.Add("-part_col_prog",&col_x.code,"color","program to compute particle color (r=...;g=...;b=...;)");
    parse_args.Add("-part_size",&part_size,"width","particle radius");
    parse_args.Add("-part_v_col",&col_v.col,"color","fixed V color");
    parse_args.Add("-part_v_col_prog",&col_v.code,"color","program to compute V color (r=...;g=...;b=...;)");
    parse_args.Add("-part_v_width",&part_v_width,"width","V line width");
    parse_args.Add("-part_v_scale",&part_v_scale,"scale","scale particle V length");
    parse_args.Add("-part_v_min_len",&part_v_min_len,"min_len","min_len particle V");
    parse_args.Add("-part_v_arrow",&part_v_arrow,"draw V as arrow");
    parse_args.Add("-part_file",&part_file,"file","particle filename");

    parse_args.Add("-clip",&clip_domain,&custom_clip,"domain","clipping domain");
    parse_args.Add("-clamp",&clamp,"clamp","clamp velocities to avoid overflows");

    parse_args.Add("-first",&first_frame,"frame","first frame to render");
    parse_args.Add("-last",&last_frame,"frame","last frame to render");
    parse_args.Add("-stride",&stride,"frame","render every n-th frame");


    
    parse_args.Parse();

    PHYSBAM_ASSERT(sim.size());
    VIEWER_DIR viewer_dir(sim);
    GRID<TV> grid;
    if(use_grid || use_u)
    {
        Read_From_File(viewer_dir.Common_File("grid"),grid);
        if(!custom_bounds) bounds=grid.Cell_Indices();
    }

    if(!custom_clip)
    {
        if(use_grid) clip_domain=grid.domain;
        else clip_domain=RANGE<TV>::Centered_Box();
    }
    
    fs::create_directories(sim+"/tex");

    col_u.Init();
    col_x.Init();
    col_v.Init();

    int out_num=0;
    for(viewer_dir.Set(0);viewer_dir.Valid_Frame(true);viewer_dir.Advance_Directory(0))
    {
        int frame=viewer_dir.frame_stack(0);
        if(frame<first_frame) continue;
        if(frame>last_frame) break;
        if((frame-first_frame)%stride) continue;

        TEX_FILE<T> tex(LOG::sprintf("%s/tex/frame-%04d.tex",sim,out_num++));
        tex.cur_format.line_join=1;

        auto draw_arrow=[&tex](const COL_PROG& prog,const TV& x,const TV& u,T width,bool arrow,T scale)
        {
            tex.cur_format.fill_style=0;
            tex.cur_format.line_style=1;
            tex.cur_format.line_color=prog.Color(x,u);
            tex.cur_format.line_width=width;
            tex.cur_format.arrow_style=arrow?"-latex":nullptr;
            tex.Draw_Object(x,x+u*scale);
        };

        if(use_grid)
        {
            tex.cur_format.line_color=grid_col;
            tex.cur_format.line_width=grid_width;
            
            for(int i=bounds.min_corner.x;i<=bounds.max_corner.x;i++)
            {
                TV A=grid.Node_X({i,bounds.min_corner.y});
                TV B=grid.Node_X({i,bounds.max_corner.y});
                tex.Draw_Object(A,B);
            }

            for(int i=bounds.min_corner.y;i<=bounds.max_corner.y;i++)
            {
                TV A=grid.Node_X({bounds.min_corner.x,i});
                TV B=grid.Node_X({bounds.max_corner.x,i});
                tex.Draw_Object(A,B);
            }
        }
        if(use_u)
        {
            ARRAY<TV,IV> u;
            Read_From_File(viewer_dir.Frame_File(u_file),u);

            RANGE<IV> range=bounds;
            range.max_corner+=1;
            range=range.Intersect(u.domain);
            for(RANGE_ITERATOR<TV::m> it(range);it.Valid();it.Next())
            {
                TV v=u(it.index);
                TV x=grid.Node_X(it.index);
                if(v.Magnitude()>clamp) v*=clamp/v.Magnitude();
                if(v.Magnitude()>u_min_len)
                    draw_arrow(col_u,x,v,u_width,u_arrow,u_scale);
            }
        }
        if(use_part)
        {
            GEOMETRY_PARTICLES<TV> particles;
            Read_From_File(viewer_dir.Frame_File(part_file),particles);

            ARRAY_VIEW<TV>* V=particles.Get_Array<TV>("V");
            
            for(int p=0;p<particles.X.m;p++)
            {
                TV v=V?(*V)(p):TV();
                TV x=particles.X(p);
                if(v.Magnitude()>clamp) v*=clamp/v.Magnitude();
                if(part_v_scale && v.Magnitude()>part_v_min_len) draw_arrow(col_v,x,v,part_v_width,part_v_arrow,part_v_scale);
                if(part_size)
                {
                    tex.cur_format.fill_style=1;
                    tex.cur_format.line_style=0;
                    tex.cur_format.fill_color=col_x.Color(x,v);
                    tex.cur_format.arrow_style=nullptr;
                    tex.Draw_Object(x,part_size);
                }
            }
        }

        tex.bounding_box=clip_domain;
    }

    return 0;
}
