//#####################################################################
// Copyright 2012.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Analytic_Tests/ANALYTIC_PROGRAM_HELPER.h>
#include <Core/Arrays/ARRAY.h>
#include <Core/Arrays/ARRAY_VIEW.h>
#include <Core/Symbolics/PROGRAM.h>
#include <Core/Symbolics/PROGRAM_COMPILER.h>
#include <Core/Symbolics/PROGRAM_CONTEXT.h>

namespace PhysBAM{
template<class T> ANALYTIC_PROGRAM_HELPER<T>::
ANALYTIC_PROGRAM_HELPER(const std::string& str,const std::string& out_name,
    int dim,int m,int n,bool want_ddx)
    :dim(dim)
{
    PHYSBAM_ASSERT(dim<=3);
    PROGRAM_COMPILER<T> pc;
    const char* axes[]={"x","y","z"};
    idx_f.Resize(dim);
    pc.Add_Input("X",dim);
    ARRAY<MEM_IN> X(dim);
    for(int i=0;i<dim;i++){
        X(i)=MEM_IN(i);
        pc.Add_Input(axes[i],MEM_IN(i));}
    MEM_IN t=pc.Add_Input("t");

    int dofs=0;
    if(m<0)
    {
        pc.Add_Output(out_name);
        dofs=1;
    }
    else if(n<0)
    {
        dofs=m;
        pc.Add_Output(out_name,m);
    }
    else
    {
        dofs=m*n;
        pc.Add_Output(out_name,m,n);
    }
    idx_f.Resize(dofs);
    for(int i=0;i<dofs;i++) idx_f(i)=MEM_OUT(i);
    pc.Parse(str);
    pc.Optimize();

    PROGRAM_COMPILER<T> pc_dt(pc),pc_dx(pc);
    idx_dt=pc_dt.Make_Output(pc_dt.Diff(idx_f,t));
    pc_dt.Optimize();

    if(dofs>1)
    {
        for(int i=0;i<dim;i++)
            idx_dx.Append(pc_dx.Make_Output(pc_dx.Diff(idx_f,X(i))));
    }
    else
    {
        auto grad=pc.Gradient(MEM_OUT{0});
        idx_dx.Resize(dim);
        for(int i=0;i<dim;i++)
            idx_dx(i)=pc.Make_Output(grad(MEM_IN{i}));
    }
    pc_dx.Optimize();

    if(want_ddx)
    {
        PROGRAM_COMPILER<T> pc_ddx(pc_dx);
        for(int i=0;i<dim;i++)
            idx_ddx.Append(pc_ddx.Make_Output(pc_ddx.Diff(idx_dx,X(i))));
        pc_ddx.Optimize();
        pc_ddx.Finalize(prog_ddx);
        context_ddx.Initialize(prog_ddx);
    }

    pc.Finalize(prog);
    pc_dt.Finalize(prog_dt);
    pc_dx.Finalize(prog_dx);
    context.Initialize(prog);
    context_dt.Initialize(prog_dt);
    context_dx.Initialize(prog_dx);
}
//#####################################################################
// Function f
//#####################################################################
template<class T> void ANALYTIC_PROGRAM_HELPER<T>::
f(ARRAY_VIEW<const T> X,T t) const
{
    assert(X.m==dim);
    for(int i=0;i<dim;i++) context.data_in(MEM_IN(i))=X(i);
    context.data_in(MEM_IN(dim))=t;
    prog.Execute(context);
}
//#####################################################################
// Function dt
//#####################################################################
template<class T> void ANALYTIC_PROGRAM_HELPER<T>::
dt(ARRAY_VIEW<const T> X,T t) const
{
    assert(X.m==dim);
    for(int i=0;i<dim;i++) context_dt.data_in(MEM_IN(i))=X(i);
    context_dt.data_in(MEM_IN(dim))=t;
    prog_dt.Execute(context_dt);
}
//#####################################################################
// Function dX
//#####################################################################
template<class T> void ANALYTIC_PROGRAM_HELPER<T>::
dX(ARRAY_VIEW<const T> X,T t) const
{
    assert(X.m==dim);
    for(int i=0;i<dim;i++) context_dx.data_in(MEM_IN(i))=X(i);
    context_dx.data_in(MEM_IN(dim))=t;
    prog_dx.Execute(context_dx);
}
//#####################################################################
// Function ddX
//#####################################################################
template<class T> void ANALYTIC_PROGRAM_HELPER<T>::
ddX(ARRAY_VIEW<const T> X,T t) const
{
    assert(X.m==dim);
    for(int i=0;i<dim;i++) context_ddx.data_in(MEM_IN(i))=X(i);
    context_ddx.data_in(MEM_IN(dim))=t;
    prog_ddx.Execute(context_ddx);
}
template class ANALYTIC_PROGRAM_HELPER<float>;
template class ANALYTIC_PROGRAM_HELPER<double>;
}
