//#####################################################################
// Copyright 2012.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYREGHT.txt.
//#####################################################################
// Class PROGRAM_COMPILER
//#####################################################################
#ifndef __PROGRAM_COMPILER__
#define __PROGRAM_COMPILER__

#include <Core/Data_Structures/HASHTABLE.h>
#include <Core/Symbolics/CODE_BLOCK.h>
#include <Core/Symbolics/COMPILER_DATA.h>
#include <Core/Symbolics/REG.h>
#include <Core/Symbolics/REG_ARRAY.h>
#include <functional>
#include <list>
#include <map>
#include <string>
namespace PhysBAM{

// NOTE: See the README file for instructions on how to use this class.

struct PROGRAM_PARSE_NODE;
template<class T> struct DISTRIBUTE_HELPER;
template<class T> struct PROGRAM;
struct INSTRUCTION;
struct OPERATION;
struct LIVE_INTERVAL;

template<class T> struct COMPILER_DATA;

template<class T>
struct PROGRAM_COMPILER
{
private:
    typedef COMPILER_DATA<T> DATA;
    typedef CODE_BLOCK::ITERATOR INST_IT;
    
    // block id, const iterator
    MEM_IN num_in{};
    MEM_OUT num_out{};
    MEM_REG num_reg{};
    ARRAY<T,MEM_CONST> constants; // starts with 0 and 1.
    std::map<T,MEM_CONST> constant_lookup;
    ARRAY<CODE_BLOCK*,MEM_LABEL> code_blocks;
    HASHTABLE<std::string,DATA*> dict;
    REG_ARRAY<INST_IT> defs;
    REG_ARRAY<HASHTABLE<const INSTRUCTION*,INST_IT>> uses;
    bool use_def_valid=false;
    HASHTABLE<std::string,std::pair<DATA(*)(PROGRAM_COMPILER<T>*pc,const ARRAY<DATA>& in),int>>& Func_Lookup();

    HASHTABLE<std::pair<REG,REG>,REG_AUG> diffs;

    struct ADJOINT
    {
        REG_AUG val[3];
        friend std::ostream& operator<<(std::ostream& o,const ADJOINT& a)
        {
            o<<'('<<a.val[0]<<' '<<a.val[1]<<' '<<a.val[2]<<')';
            return o;
        }
    };

    HASHTABLE<REG,ADJOINT> adjoints;
    HASHTABLE<REG,INST_IT> diff_block_end;
    HASHTABLE<REG,REG> diff_parent;
    ARRAY<MEM_LABEL> reversal_points;

    struct TRACE_REG
    {
        ARRAY<std::pair<MEM_LABEL,REG>> mapping;
    };
    // Makes register available until end of program
    HASHTABLE<REG,REG> parent_reg;
    HASHTABLE<REG,TRACE_REG> trace_reg;

    ARRAY<std::pair<std::string,DATA>> var_in;
    ARRAY<std::pair<std::string,DATA>> var_out;

public:

    bool debug=false;
    bool verbose=false;

    struct VARIABLE_STATE
    {
        T value;
        enum state_type {unknown,constant,overdetermined} state;
    };

    PROGRAM_COMPILER();
    PROGRAM_COMPILER(const PROGRAM_COMPILER& pc);
    ~PROGRAM_COMPILER();

    // Add scalar.
    MEM_IN Add_Input(const std::string& name,MEM_IN in=MEM_IN(-1));

    // Add vector.  Returns first register.  Remainder are consecutive.
    MEM_IN Add_Input(const std::string& name,int m,MEM_IN in=MEM_IN(-1));

    // Add matrix.  Returns first register.  Remainder are consecutive.  Column-major order.
    MEM_IN Add_Input(const std::string& name,int m,int n,MEM_IN in=MEM_IN(-1));

    // Add scalar.
    MEM_OUT Add_Output(const std::string& name,MEM_OUT out=MEM_OUT(-1));

    // Add vector.  Returns first register.  Remainder are consecutive.
    MEM_OUT Add_Output(const std::string& name,int m,MEM_OUT out=MEM_OUT(-1));

    // Add matrix.  Returns first register.  Remainder are consecutive.  Column-major order.
    MEM_OUT Add_Output(const std::string& name,int m,int n,MEM_OUT out=MEM_OUT(-1));

    // Differentiation; Note that these routines do not create new output
    // variables, but you can do so with Make_Output.
    REG Diff(REG func_var,REG diff_var);
    ARRAY<REG,MEM_OUT> Diff(REG diff_var); // func_var is all outputs
    ARRAY<REG> Diff(ARRAY_VIEW<const REG> func_vars,REG diff_var);
    template<class R,class T_ARRAY,class I>
    ARRAY<REG,I> Diff(const ARRAY_BASE<R,T_ARRAY,I>& func_vars,REG diff_var)
    {
        if constexpr(is_same<R,REG>::value && HAS_POINTER<T_ARRAY>::value)
            return ARRAY<REG,I>(Diff(ARRAY_VIEW<const REG>(func_vars,force_idx),diff_var),force_idx);
        else
        {
            ARRAY<REG> fv(func_vars,cast_data,force_idx);
            return ARRAY<REG,I>(Diff(ARRAY_VIEW<const REG>(fv),diff_var),force_idx);
        }
    }
    ARRAY<REG> Diff(ARRAY_VIEW<const MEM_OUT> func_vars,REG diff_var);
    ARRAY<REG,MEM_IN> Gradient(REG func_var); // diff_var is all inputs
    // ret(out_var)(in_var)
    ARRAY<ARRAY<REG,MEM_IN>> Gradient(ARRAY_VIEW<const REG> func_vars); // diff_var is all inputs
    REG Directional_Derivative(REG func_var,ARRAY_VIEW<const REG,MEM_IN> dir);
    ARRAY<REG,MEM_OUT> Directional_Derivative(ARRAY_VIEW<const REG,MEM_IN> dir); // func_var is all outputs
    
    void Parse(const std::string& str);
    void Print() const;
    void Optimize();
    MEM_CONST Add_Constant(T x);
    void Finalize(PROGRAM<T>& prog);

    void Shift_Registers();
    void Rename_Variable(REG var_old,REG var_new);

    MEM_OUT Append(const PROGRAM_COMPILER<T>& pc);
    MEM_REG Num_Reg() const {return num_reg;}
    MEM_OUT Num_Out() const {return num_out;}
    MEM_IN Num_In() const {return num_in;}
    MEM_REG New_Reg() {return {num_reg++};}
    MEM_OUT New_Out() {return {num_out++};}
    
    // Makes new output, copies var to it.
    MEM_OUT Make_Output(REG var);
    template<class A,class I>
    auto Make_Output(const ARRAY<A,I>& var)
    {
        ARRAY<decltype(Make_Output(var(I{}))),I> a(var.m);
        for(I i{0};i<var.m;i++) a(i)=Make_Output(var(i));
        return a;
    }
    void Test_Invariants() const;

    // Turns outputs into regular registers.  Returns first such register and
    // the number of former outputs
    std::pair<MEM_REG,int> Move_Outputs_To_Reg();

    void Register_Default_Func();
    void Register_Func(const std::string& name,int num_in,DATA(*func)(PROGRAM_COMPILER<T>*pc,const ARRAY<DATA>& in));
    T Get_Const(const DATA& a) const;

private:
    bool Deduce_Op(OPERATION op,T& out,T* in0,T* in1,T* in2) const;
    DATA Exec_Func(const std::string& name,const ARRAY<DATA>& in);

    INST_IT Begin_It() const;
    INST_IT Last_It() const;
    INST_IT End_It() const;
    void Move_Past_Phis(INST_IT& it) const;
    void Print(const CODE_BLOCK* B) const;
    INST_IT Get_Def(REG var) const;
    void Set_Def(REG var,INST_IT def);
    const HASHTABLE<const INSTRUCTION*,INST_IT>& Get_Uses(REG var) const;
    HASHTABLE<const INSTRUCTION*,INST_IT>& Get_Uses(REG var);
    const HASHTABLE<const INSTRUCTION*,INST_IT>& Get_Uses(const INSTRUCTION& in) const;
    void Remove_Use_Def(INST_IT it);
    void Add_Use_Def(INST_IT it);
    void Delete_Instruction(INST_IT it);
    void Parse_Command(const std::string& str);
    void Print_Node(PROGRAM_PARSE_NODE* node);
    DATA Process_Node(PROGRAM_PARSE_NODE* node);
    DATA& Process_Assign(DATA& a,const DATA& b);
public:
    DATA Process_Componentwise(const DATA& a,const DATA& b,const OPERATION& op,
        bool allow_s_mv,const char* s0,const char* s1,const char* s2);
    DATA Process_Componentwise(const DATA& a,const OPERATION& op);
    DATA Process_Mod(const DATA& a,const DATA& b,bool allow_s_mv,
        const char* s0,const char* s1,const char* s2);
    DATA Process_Mul(const DATA& a,const DATA& b);
private:

    ADJOINT Make_Adjoint(REG func_var);
    REG_AUG Diff_Internal(REG func_var,REG diff_var);
    void Make_SSA();
    void Make_SSA_Relabel(MEM_REG& count,ARRAY<ARRAY<REG>,MEM_REG>& S,HASHTABLE<REG>& V_used,const ARRAY<ARRAY<MEM_LABEL>,MEM_LABEL>& children,MEM_LABEL bl);
    void Leave_SSA();
    void Relabel_Registers(ARRAY<REG,MEM_REG>& var_map);
    void Compute_Use_Def();
    void Invalidate_Use_Def();
    void Copy_Propagation();
    void Simplify_Phis();
    void Sparse_Conditional_Constant_Propagation();
    void SCCP_Visit_Instruction(INST_IT it,REG_ARRAY<VARIABLE_STATE>& variable_state,
        ARRAY<CODE_BLOCK*>& block_worklist,ARRAY<INST_IT>& op_worklist,ARRAY<bool,MEM_LABEL>& block_exec);
    void Remove_Dead_Code();
    void Detach_Next(CODE_BLOCK* B,int j);
    void Propagate_Copy(REG old_var,REG new_var);
    void Local_Common_Expresssion_Elimination(CODE_BLOCK* B);
    void Local_Common_Expresssion_Elimination();
    void Reduce_In_Place(INST_IT it);
    void Compress_Registers();
    void Eliminate_Unused_Register(MEM_REG var);
    void Eliminate_Unused_Constant(MEM_CONST var);

    // Inserts instructions before it.  Allocates output and returns them.
    REG Before_Instruction(const INST_IT& it,const OPERATION& type,REG src0,REG src1={},REG src2={});

public:
    REG Append_Instruction(const OPERATION& type,REG src0,REG src1={},REG src2={})
    {return Before_Instruction(End_It(),type,src0,src1,src2);}

    // make_true and make_false should add instructions (with
    // Append_Instruction) to compute quantities.  The returned arrays must be
    // the same length.  By emitting instructions on demand, this routine allows
    // lazy execution of that code in the resulting program.
    ARRAY<REG> Append_If_Else(REG cond,std::function<ARRAY<REG>()> make_true,std::function<ARRAY<REG>()> make_false);
private:
    REG Append_Instruction(CODE_BLOCK* B,const OPERATION& type,REG src0,REG src1={},REG src2={});
    REG Prepend_Instruction(CODE_BLOCK* B,const OPERATION& type,REG src0,REG src1={},REG src2={});
    REG Prepend_Phi(CODE_BLOCK* B,REG src0,REG src1);
    REG_AUG Prepend_Phi(CODE_BLOCK* B,REG_AUG src0,REG_AUG src1);

    REG Append_Instruction_Dest(CODE_BLOCK* B,const OPERATION& type,REG dest,REG src0,REG src1={},REG src2={});
    REG Append_Instruction_Dest(const OPERATION& type,REG dest,REG src0,REG src1={},REG src2={});


    
    void Propagate_Copy_And_Remove_Instruction(INST_IT it);
    void Simplify_With_Distributive_Law_Helper_Prod(REG var,ARRAY<INST_IT>& extra_nodes,DISTRIBUTE_HELPER<T>& prod);
    void Simplify_With_Distributive_Law_Helper(REG var,T coeff,ARRAY<INST_IT>& extra_nodes,ARRAY<DISTRIBUTE_HELPER<T> >& summands);
    REG Simplify_With_Distributive_Law_Emit_Prod(DISTRIBUTE_HELPER<T>& prod,INST_IT before,bool& need_neg);
    REG Simplify_With_Distributive_Law_Emit_Simple(ARRAY<DISTRIBUTE_HELPER<T> >& summands,INST_IT before,bool& need_neg);
    REG Simplify_With_Distributive_Law_Emit_Consts(ARRAY<DISTRIBUTE_HELPER<T> >& summands,INST_IT before,bool& need_neg);
    REG Simplify_With_Distributive_Law_Emit(ARRAY<DISTRIBUTE_HELPER<T> >& summands,INST_IT before,bool& need_neg);
    INST_IT Simplify_With_Distributive_Law(INST_IT it);
    void Simplify_With_Distributive_Law();
    void Change_Instruction(INST_IT it,const INSTRUCTION& in);
    void Ensure_Valid_Use_Def(REG var);
    void Merge_Keep(CODE_BLOCK* B,std::map<REG,REG>& r,
        const std::map<REG,REG>& a,const std::map<REG,REG>& b);

    // Insert instructions before it if necessary so that ra no longer has
    // neg/inv bit set; does not logically change value of input.
    REG_AUG Clear_Neg(REG_AUG ra);
    REG_AUG Clear_Inv(REG_AUG ra);
    REG Clear_Aug(REG_AUG ra);
    void Make_Same_Neg(REG_AUG& ra,REG_AUG& rb);
    void Make_Same_Inv(REG_AUG& ra,REG_AUG& rb);
    void Make_Same_Aug(REG_AUG& ra,REG_AUG& rb);

    // Logically negate input without inserting instructions
    REG_AUG Neg(REG_AUG ra){ra.neg=!ra.neg;return ra;}
    REG_AUG Inv(REG_AUG ra){ra.inv=!ra.inv;return ra;}

    // Emit just enough instructions before it to perform the logical operation.
    REG_AUG Mul(INST_IT it,REG_AUG a,REG_AUG b);
    REG_AUG Add(INST_IT it,REG_AUG a,REG_AUG b);
    void Add_Mul(INST_IT it,REG_AUG& o,REG_AUG a,REG_AUG b);
    REG_AUG If_Aug(INST_IT it,REG a,REG_AUG b,REG_AUG c);
    void Clear_Diff_Structures();
    void Add_Reversal_Point();
    MEM_LABEL Get_Reversal_Point(MEM_LABEL id);
    REG Trace_Register(REG reg,MEM_LABEL id);
    REG_AUG Trace_Register(REG_AUG reg,MEM_LABEL id);
    void Compute_Live(HASHTABLE<REG,ARRAY<LIVE_INTERVAL>>& live_ranges) const;
};
}
#endif
