//#####################################################################
// Copyright 2012-2025 Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYREGHT.txt.
//#####################################################################
#include <Core/Symbolics/COMPILER_DATA_OPS.h>
#include <Core/Symbolics/COMPILER_DATA.h>
#include <Core/Symbolics/PROGRAM_COMPILER.h>
namespace PhysBAM{
#define REG_F1(f)                                                       \
    template<class T> auto COMPILER_DATA_OPS<T>::f(PC*pc,const ARRAY<D>& in) -> D \
    {return COMPILER_DATA_OPS<T>{pc}.f(in(0));};                        \
    template<class T> auto COMPILER_DATA_OPS<T>::f(const D& a) -> D
#define REG_F2(f)                                                       \
    template<class T> auto COMPILER_DATA_OPS<T>::f(PC*pc,const ARRAY<D>& in) -> D \
    {return COMPILER_DATA_OPS<T>{pc}.f(in(0),in(1));};                  \
    template<class T> auto COMPILER_DATA_OPS<T>::f(const D& a,const D& b) -> D
#define REG_F(f)                                                        \
    template<class T> auto COMPILER_DATA_OPS<T>::f(PC*pc,const ARRAY<D>& in) -> D
#define R1(op) REG_F1(op){return pc->Process_Componentwise(a,op_##op);}
#define R2(op) REG_F2(op){return pc->Process_Componentwise(a,b,op_##op,true,#op "(",",",")");}
R1(inv);
R1(sqrt);
R1(exp);
R1(ln);
R1(sin);
R1(cos);
R1(asin);
R1(acos);
R1(atan);
R2(atan2);
R1(floor);
R1(ceil);
R1(rint);
R1(abs);
R2(min);
R2(max);
R2(pow);
R1(neg);
R2(add);
R2(sub);
R2(mul);
R2(div);
REG_F2(mod){return pc->Process_Mod(a,b,true,"mod(",",",")");}
REG_F1(tan){return div(sin(a),cos(a));}
REG_F1(cot){return div(cos(a),sin(a));}
REG_F1(sec){return inv(cos(a));}
REG_F1(csc){return inv(sin(a));}
REG_F1(acot){return atan(inv(a));}
REG_F1(asec){return acos(inv(a));}
REG_F1(acsc){return asin(inv(a));}
REG_F1(trans)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::mat);
    COMPILER_DATA<T> r(a.n,a.m);
    for(int i=0;i<a.m;i++)
        for(int j=0;j<a.n;j++)
            r(j,i)=a(i,j);
    return r;
}
REG_F2(dot)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::vec);
    PHYSBAM_ASSERT(b.type==COMPILER_DATA<T>::TYPE::vec);
    PHYSBAM_ASSERT(a.m==b.m);
    REG x=MEM_CONST(0);
    for(int i=0;i<a.m;i++)
        x=pc->Append_Instruction(op_add,x,pc->Append_Instruction(op_mul,a.x(i),b.x(i)));
    return COMPILER_DATA<T>(x);
}
REG_F2(cross)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::vec);
    PHYSBAM_ASSERT(b.type==COMPILER_DATA<T>::TYPE::vec);
    if(a.m+b.m==1)
    {
        COMPILER_DATA<T> r(1);
        r.x(0)=MEM_CONST(0);
        return r;
    }
    if(a.m==1 && b.m==1) return COMPILER_DATA<T>(0);
    if(a.m==1 && b.m==2)
    {
        COMPILER_DATA<T> r(2);
        r.x(0)=pc->Append_Instruction(op_neg,
            pc->Append_Instruction(op_mul,a.x(0),b.x(1)));
        r.x(1)=pc->Append_Instruction(op_mul,a.x(0),b.x(0));
        return r;
    }
    if(a.m==2 && b.m==1)
    {
        COMPILER_DATA<T> r(2);
        r.x(0)=pc->Append_Instruction(op_mul,a.x(1),b.x(0));
        r.x(1)=pc->Append_Instruction(op_neg,
            pc->Append_Instruction(op_mul,a.x(0),b.x(0)));
        return r;
    }
    if(a.m==2 && b.m==2)
    {
        COMPILER_DATA<T> r(1);
        r.x(0)=pc->Append_Instruction(op_sub,
            pc->Append_Instruction(op_mul,a.x(0),b.x(1)),
            pc->Append_Instruction(op_mul,a.x(1),b.x(0)));
        return r;
    }
    PHYSBAM_ASSERT(a.m==3 && b.m==3);
    COMPILER_DATA<T> r(3);
    r.x(0)=pc->Append_Instruction(op_sub,
        pc->Append_Instruction(op_mul,a.x(1),b.x(2)),
        pc->Append_Instruction(op_mul,a.x(2),b.x(1)));
    r.x(1)=pc->Append_Instruction(op_sub,
        pc->Append_Instruction(op_mul,a.x(2),b.x(0)),
        pc->Append_Instruction(op_mul,a.x(0),b.x(2)));
    r.x(2)=pc->Append_Instruction(op_sub,
        pc->Append_Instruction(op_mul,a.x(0),b.x(1)),
        pc->Append_Instruction(op_mul,a.x(1),b.x(0)));
    return r;
}
REG_F2(outer)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::vec);
    PHYSBAM_ASSERT(b.type==COMPILER_DATA<T>::TYPE::vec);
    COMPILER_DATA<T> r(a.m,b.m);
    for(int i=0;i<a.m;i++)
        for(int j=0;j<b.m;j++)
            r(i,j)=pc->Append_Instruction(op_mul,a.x(i),b.x(j));
    return r;
}
REG_F1(tr)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::mat);
    PHYSBAM_ASSERT(a.m==a.n);
    REG x=MEM_CONST(0);
    for(int i=0;i<a.m;i++)
        x=pc->Append_Instruction(op_add,x,a(i,i));
    return COMPILER_DATA<T>(x);
}
REG_F(zeros)
{
    PHYSBAM_ASSERT(in.m==1 || in.m==2);
    if(in.m==1)
    {
        COMPILER_DATA<T> r(pc->Get_Const(in(0)));
        r.x.Fill(MEM_CONST(0));
        return r;
    }
    COMPILER_DATA<T> r(pc->Get_Const(in(0)),pc->Get_Const(in(1)));
    r.x.Fill(MEM_CONST(0));
    return r;
}
REG_F(ones)
{
    PHYSBAM_ASSERT(in.m==1 || in.m==2);
    if(in.m==1)
    {
        COMPILER_DATA<T> r(pc->Get_Const(in(0)));
        r.x.Fill(MEM_CONST(1));
        return r;
    }
    COMPILER_DATA<T> r(pc->Get_Const(in(0)),pc->Get_Const(in(1)));
    r.x.Fill(MEM_CONST(1));
    return r;
}
REG_F1(id)
{
    int m=pc->Get_Const(a);
    COMPILER_DATA<T> r(m,m);
    r.x.Fill(MEM_CONST(0));
    for(int i=0;i<m;i++) r(i,i)=MEM_CONST(1);
    return r;
}
REG_F1(norm2)
{
    REG x=MEM_CONST(0);
    for(auto&a:a.x)
        x=pc->Append_Instruction(op_add,x,pc->Append_Instruction(op_mul,a,a));
    return COMPILER_DATA<T>(x);
}
REG_F1(norm)
{
    if(a.x.m==1)
        return COMPILER_DATA<T>(pc->Append_Instruction(op_abs,a.x(0)));
    REG x=MEM_CONST(0);
    for(auto&a:a.x)
        x=pc->Append_Instruction(op_add,x,pc->Append_Instruction(op_mul,a,a));
    return COMPILER_DATA<T>(pc->Append_Instruction(op_sqrt,x));
}
REG_F1(det)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::mat);
    PHYSBAM_ASSERT(a.m==a.n);
    if(a.m==0) return COMPILER_DATA<T>(MEM_CONST(0));
    if(a.m==1) return COMPILER_DATA<T>(a.x(0));
    if(a.m==2) return sub(mul(a.x(0),a.x(3)),mul(a.x(1),a.x(2)));
    PHYSBAM_ASSERT(a.m==3);
    return dot(cross(col(a,0),col(a,1)),col(a,2));
}
REG_F(vec)
{
    COMPILER_DATA<T> r(0);
    for(auto&i:in)
    {
        PHYSBAM_ASSERT(i.type==COMPILER_DATA<T>::TYPE::scalar || i.type==COMPILER_DATA<T>::TYPE::vec);
        r.x.Append(i.x);
    }
    r.m=r.x.m;
    return r;
}
REG_F2(row)
{
    return row(a,pc->Get_Const(b));
}
template<class T> auto COMPILER_DATA_OPS<T>::row(const D& a,int k) -> D
{
    COMPILER_DATA<T> r(a.n);
    for(int i=0;i<a.n;i++) r.x(i)=a(k,i);
    return r;
}
REG_F2(col)
{
    return col(a,pc->Get_Const(b));
}
template<class T> auto COMPILER_DATA_OPS<T>::col(const D& a,int k) -> D
{
    COMPILER_DATA<T> r(a.m);
    for(int i=0;i<a.m;i++) r.x(i)=a(i,k);
    return r;
}
REG_F1(inverse)
{
    PHYSBAM_ASSERT(a.type==COMPILER_DATA<T>::TYPE::mat);
    PHYSBAM_ASSERT(a.m==a.n);
    if(a.m==1) return inv(a);
    if(a.m==2)
    {
        COMPILER_DATA<T> r(2,2);
        r(0,0)=a(1,1);
        r(1,1)=a(0,0);
        r(0,1)=neg(a(0,1)).x(0);
        r(1,0)=neg(a(1,0)).x(0);
        return div(r,det(a));
    }
    PHYSBAM_ASSERT(a.m==3);
    auto u=row(a,0),v=row(a,1),w=row(a,2);
    auto uv=cross(u,v),vw=cross(v,w),wu=cross(w,u);
    COMPILER_DATA<T> r(3,3);
    for(int i=0;i<3;i++) r(i,0)=vw.x(i);
    for(int i=0;i<3;i++) r(i,1)=wu.x(i);
    for(int i=0;i<3;i++) r(i,2)=uv.x(i);
    return div(r,dot(uv,w));
}
template struct COMPILER_DATA_OPS<float>;
template struct COMPILER_DATA_OPS<double>;
}
