//#####################################################################
// Copyright 2004-2009, Eran Guendelman, Frank Losasso, 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/Symbolics/PARTICLE_ATTRIBUTE_PROGRAM.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Core/Matrices/FRAME.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
#include <Core/Particles/PARTICLES.h>
#include <Core/Symbolics/PROGRAM.h>
#include <Core/Symbolics/PROGRAM_COMPILER.h>
#include <Core/Symbolics/PROGRAM_CONTEXT.h>
#include <Core/Vectors/TWIST.h>
#include <Core/Vectors/VECTOR.h>
#include <regex>
#include <stdexcept>
namespace PhysBAM{
template<class TV> class PARTICLE_ATTRIBUTE_PROGRAM_HELPER;

namespace{
enum TYPE
{
    type_scalar,type_vector,
    type_matrix,type_sym,type_diag,
    type_frame_t,type_frame_r,type_twist_l,type_twist_a
};
}
template<class TV>
struct ATTR_BASE
{
    virtual ~ATTR_BASE()=default;
    virtual bool Fetch_Attribute()=0;
    virtual void Read_Attribute(int p) const=0;
    virtual void Write_Attribute(int p)=0;
};

template<class TV>
class PARTICLE_ATTRIBUTE_PROGRAM_HELPER
{
public:
    typedef typename TV::SCALAR T;
    std::string code;
    PARTICLES<TV>* particles=nullptr;
    ARRAY<std::string> out_attrs;
    PROGRAM_COMPILER<T> pc;
    PROGRAM<T> prog;
    PROGRAM_CONTEXT<T> context;
    ARRAY<ATTR_BASE<TV>*> attrs;
};

template<class OBJ,int type,class TV>
struct ATTR:ATTR_BASE<TV>
{
    typedef typename TV::SCALAR T;
    PARTICLE_ATTRIBUTE_PROGRAM_HELPER<TV>* h;
    std::string attr_name;
    MEM_IN in;
    MEM_OUT out;
    bool is_out;
    ARRAY_VIEW<OBJ>* array=nullptr;

    ATTR(PARTICLE_ATTRIBUTE_PROGRAM_HELPER<TV>* h,const std::string& attr_name,
        const std::string& var_name,int m,int n)
        :h(h),attr_name(attr_name),is_out(h->out_attrs.Contains(attr_name))
    {
        if(is_out)
        {
            if(m<=0) out=h->pc.Add_Output(var_name);
            else if(n<=0) out=h->pc.Add_Output(var_name,m);
            else out=h->pc.Add_Output(var_name,m,n);
        }
        else
        {
            if(m<=0) in=h->pc.Add_Input(var_name);
            else if(n<=0) in=h->pc.Add_Input(var_name,m);
            else in=h->pc.Add_Input(var_name,m,n);
        }
    }
    bool Fetch_Attribute() override
    {
        array=h->particles->template Get_Array<OBJ>(attr_name);
        return array!=nullptr;
    }
    void Read_Attribute(int p) const override
    {
        if(is_out) return;
        const OBJ& x=(*array)(p);
        if constexpr(type==type_scalar)
            h->context.Set(in,(T)x);
        else if constexpr(type==type_vector)
            h->context.Set(in,VECTOR<T,OBJ::m>(x));
        else if constexpr(type==type_matrix)
            h->context.Set(in,MATRIX<T,OBJ::m,OBJ::n>(x));
        else if constexpr(type==type_diag)
            h->context.Set(in,MATRIX<T,OBJ::m>(DIAGONAL_MATRIX<T,OBJ::m>(x)));
        else if constexpr(type==type_sym)
            h->context.Set(in,MATRIX<T,OBJ::m>(SYMMETRIC_MATRIX<T,OBJ::m>(x)));
        else if constexpr(type==type_frame_t)
            h->context.Set(in,VECTOR<T,OBJ::VECTOR_T::m>(x.t));
        else if constexpr(type==type_frame_r)
        {
            auto r=x.r.Rotation_Matrix();
            h->context.Set(in,MATRIX<T,decltype(r)::m>(r));
        }
        else if constexpr(type==type_twist_l)
            h->context.Set(in,VECTOR<T,decltype(x.linear)::m>(x.linear));
        else if constexpr(type==type_twist_a)
            h->context.Set(in,VECTOR<T,decltype(x.angular)::m>(x.angular));
        else PHYSBAM_FATAL_ERROR();
    }
    void Write_Attribute(int p) override
    {
        if(!is_out) return;
        OBJ& x=(*array)(p);
        if constexpr(type==type_scalar)
            x=(OBJ)h->context.Get(out);
        else if constexpr(type==type_vector)
            x=(OBJ)h->context.template Get<OBJ::m>(out);
        else if constexpr(type==type_matrix)
            x=(OBJ)h->context.template Get<OBJ::m,OBJ::n>(out);
        else if constexpr(type==type_diag)
            x=(OBJ)h->context.template Get<OBJ::m,OBJ::n>(out).Diagonal_Part();
        else if constexpr(type==type_sym)
            x=(OBJ)h->context.template Get<OBJ::m,OBJ::n>(out).Symmetric_Part();
        else if constexpr(type==type_frame_t)
            x.t=(decltype(x.t))h->context.template Get<decltype(x.t)::m>(out);
        else if constexpr(type==type_frame_r)
        {
            typedef typename OBJ::VECTOR_T::SCALAR S;
            constexpr int d=OBJ::VECTOR_T::m;
            x.r=(decltype(x.r))(MATRIX<S,d>)h->context.template Get<d,d>(out);
        }
        else if constexpr(type==type_twist_l)
            x.linear=(decltype(x.linear))h->context.template Get<decltype(x.linear)::m>(out);
        else if constexpr(type==type_twist_a)
            x.angular=(decltype(x.angular))h->context.template Get<decltype(x.angular)::m>(out);
        else PHYSBAM_FATAL_ERROR();
    }
};

template<class TV,class U,int t>
void Test(auto*h,auto*a,const auto&s,int m,int n)
{
    if(auto e=dynamic_cast<ARRAY_COLLECTION_ELEMENT<U>*>(a))
    {
        std::regex r("\\b"+s+"\\b");
        if(!std::regex_search(h->code,r)) return;
        h->attrs.Append(new ATTR<U,t,TV>(h,a->name,s,m,n));
    }
}
template<class TV,class U,int d>
void Test_Td(auto*h,auto*a)
{
    constexpr int sp=VECTOR<U,d>::SPIN::m;
    Test<TV,VECTOR<U,d>,type_vector>(h,a,a->name,d,-1);
    Test<TV,MATRIX<U,d>,type_matrix>(h,a,a->name,d,d);
    Test<TV,SYMMETRIC_MATRIX<U,d>,type_sym>(h,a,a->name,d,d);
    Test<TV,DIAGONAL_MATRIX<U,d>,type_diag>(h,a,a->name,d,d);
    Test<TV,FRAME<VECTOR<U,d>>,type_frame_t>(h,a,a->name+"_t",d,-1);
    Test<TV,FRAME<VECTOR<U,d>>,type_frame_r>(h,a,a->name+"_r",sp,sp);
    Test<TV,TWIST<VECTOR<U,d>>,type_twist_l>(h,a,a->name+"_l",d,-1);
    Test<TV,TWIST<VECTOR<U,d>>,type_twist_a>(h,a,a->name+"_a",sp,-1);
}
template<class TV,int d>
void Test_d(auto*h,auto*a)
{
    Test<TV,VECTOR<int,d>,type_vector>(h,a,a->name,d,-1);
    Test_Td<TV,float,d>(h,a);
    Test_Td<TV,double,d>(h,a);
}
template<class TV>
void Test_All(auto*h,auto*a)
{
    Test<TV,bool,type_scalar>(h,a,a->name,-1,-1);
    Test<TV,int,type_scalar>(h,a,a->name,-1,-1);
    Test<TV,float,type_scalar>(h,a,a->name,-1,-1);
    Test<TV,double,type_scalar>(h,a,a->name,-1,-1);
    Test_d<TV,1>(h,a);
    Test_d<TV,2>(h,a);
    Test_d<TV,3>(h,a);
}
//#####################################################################
// Constructor
//#####################################################################
template<class TV> PARTICLE_ATTRIBUTE_PROGRAM<TV>::
PARTICLE_ATTRIBUTE_PROGRAM(PARTICLES<TV>& particles,const std::string& code,const ARRAY<std::string>& out_attrs)
    :h(new PARTICLE_ATTRIBUTE_PROGRAM_HELPER<TV>)
{
    h->code=code;
    h->out_attrs=out_attrs;
    h->particles=&particles;
    for(auto* a:particles.arrays)
        Test_All<TV>(h,a);
    h->pc.Parse(h->code);
    h->pc.Optimize();
    h->pc.Finalize(h->prog);
    h->context.Initialize(h->prog);
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV> PARTICLE_ATTRIBUTE_PROGRAM<TV>::
~PARTICLE_ATTRIBUTE_PROGRAM()
{
    for(auto*p:h->attrs) delete p;
    delete h;
}
//#####################################################################
// Function Compute_Attributes
//#####################################################################
template<class TV> bool PARTICLE_ATTRIBUTE_PROGRAM<TV>::
Compute_Attributes()
{
    for(auto*a:h->attrs)
        if(!a->Fetch_Attribute())
            return false;
    for(int p=0;p<h->particles->number;p++)
    {
        for(auto*a:h->attrs)
            a->Read_Attribute(p);
        h->prog.Execute(h->context);
        for(auto*a:h->attrs)
            a->Write_Attribute(p);
    }
    return true;
}
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<float,1> >;
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<float,2> >;
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<float,3> >;
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<double,1> >;
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<double,2> >;
template class PARTICLE_ATTRIBUTE_PROGRAM<VECTOR<double,3> >;
}
