//#####################################################################
// Copyright 2020, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
// Class AUTO_DIFF_DYN_REV
//##################################################################### 
#ifndef __AUTO_DIFF_DYN_REV__
#define __AUTO_DIFF_DYN_REV__

#include <Core/Math_Tools/cube.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
#include <Core/Vectors/SCALAR_POLICY.h>
#include <Core/Vectors/VECTOR.h>
#include <cmath>
using ::std::sin;
using ::std::cos;
using ::std::tan;
using ::std::exp;
using ::std::log;
using ::std::abs;
using ::std::sqrt;
using ::PhysBAM::sqr;
using ::PhysBAM::cube;
using ::std::hypot;
namespace PhysBAM{

namespace ADR{
using ::std::sin;
using ::std::cos;
using ::std::tan;
using ::std::asin;
using ::std::acos;
using ::std::atan;
using ::std::sinh;
using ::std::cosh;
using ::std::tanh;
using ::std::asinh;
using ::std::acosh;
using ::std::atanh;
using ::std::atan2;
using ::std::exp;
using ::std::log;
using ::std::abs;
using ::std::sqrt;
using ::PhysBAM::sqr;
using ::PhysBAM::cube;
using ::std::hypot;
using ::std::floor;
using ::std::ceil;
using ::std::rint;

enum adr_op_type : unsigned char
{
    op_nop,op_copy,op_add,op_sub,op_mul,op_div,op_neg,op_inv,
    op_sqrt,op_exp,op_ln,
    op_sin,op_cos,op_tan,op_asin,op_acos,op_atan,op_atan2,
    op_sinh,op_cosh,op_tanh,op_asinh,op_acosh,op_atanh,
    op_hypot,
    op_floor,op_ceil,op_rint,
    op_last
};

enum class adr_arg_type : unsigned char
{
    arg_none,arg_a,arg_x,arg_ab,arg_ax,arg_xa
};

template<class T>
struct TABLE
{
    struct ENTRY
    {
        const char* name;
        const char* pattern;
        T(*op)(T a,T b);
        void(*rev_op)(T a,T& da,T b,T& db,T r,T dr);
    };
    static constexpr ENTRY entries[op_last]=
    {
        {
            "nop","0",
            [](T a,T b){return (T)0;},
            [](T a,T& da,T b,T& db,T r,T dr){},
        },
        {
            "copy","%s",
            [](T a,T b){return a;},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr;},
        },
        {
            "+","%s + %s",
            [](T a,T b){return a+b;},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr;db+=dr;},
        },
        {
            "-","%s - %s",
            [](T a,T b){return a-b;},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr;db-=dr;},
        },
        {
            "*","%s * %s",
            [](T a,T b){return a*b;},
            [](T a,T& da,T b,T& db,T r,T dr){da+=b*dr;db+=a*dr;},
        },
        {
            "/","%s / %s",
            [](T a,T b){return a/b;},
            [](T a,T& da,T b,T& db,T r,T dr){dr/=b;da+=dr;db-=r*dr;},
        },
        {
            "-","-%s",
            [](T a,T b){return -a;},
            [](T a,T& da,T b,T& db,T r,T dr){da-=dr;},
        },
        {
            "inv","1 / %s",
            [](T a,T b){return 1/a;},
            [](T a,T& da,T b,T& db,T r,T dr){da-=r*r*dr;},
        },
        {
            "sqrt","sqrt(%s)",
            [](T a,T b){return sqrt(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/(2*r);},
        },
        {
            "exp","exp(%s)",
            [](T a,T b){return exp(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=r*dr;},
        },
        {
            "log","log(%s)",
            [](T a,T b){return log(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/a;},
        },
        {
            "sin","sin(%s)",
            [](T a,T b){return sin(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=cos(a)*dr;},
        },
        {
            "cos","cos(%s)",
            [](T a,T b){return cos(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da-=sin(a)*dr;},
        },
        {
            "tan","tan(%s)",
            [](T a,T b){return tan(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=(r*r+1)*dr;},
        },
        {
            "asin","asin(%s)",
            [](T a,T b){return asin(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/sqrt(1-a*a);},
        },
        {
            "acos","acos(%s)",
            [](T a,T b){return acos(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da-=dr/sqrt(1-a*a);},
        },
        {
            "atan","atan(%s)",
            [](T a,T b){return atan(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/(1+a*a);},
        },
        {
            "atan2","atan2(%s,%s)",
            [](T a,T b){return atan2(a,b);},
            [](T a,T& da,T b,T& db,T r,T dr){dr/=(a*a+b*b);da+=b*dr;db-=a*dr;},
        },
        {
            "sinh","sinh(%s)",
            [](T a,T b){return sinh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=cosh(a)*dr;},
        },
        {
            "cosh","cosh(%s)",
            [](T a,T b){return cosh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=sinh(a)*dr;},
        },
        {
            "tanh","tanh(%s)",
            [](T a,T b){return tanh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=(1-r*r)*dr;},
        },
        {
            "asinh","asinh(%s)",
            [](T a,T b){return asinh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/sqrt(a*a+1);},
        },
        {
            "acosh","acosh(%s)",
            [](T a,T b){return acosh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/sqrt(a*a-1);},
        },
        {
            "atanh","atanh(%s)",
            [](T a,T b){return atanh(a);},
            [](T a,T& da,T b,T& db,T r,T dr){da+=dr/(1-a*a);},
        },
        {
            "hypot","hypot(%s,%s)",
            [](T a,T b){return hypot(a,b);},
            [](T a,T& da,T b,T& db,T r,T dr){dr/=r;da+=a*dr;db+=b*dr;},
        },
        {
            "floor","floor(%s)",
            [](T a,T b){return floor(a);},
            [](T a,T& da,T b,T& db,T r,T dr){}
        },
        {
            "ceil","ceil(%s)",
            [](T a,T b){return ceil(a);},
            [](T a,T& da,T b,T& db,T r,T dr){}
        },
        {
            "rint","rint(%s)",
            [](T a,T b){return rint(a);},
            [](T a,T& da,T b,T& db,T r,T dr){}
        }
    };
};

template<class T>
struct ADR_OP
{
    adr_op_type op_type;
    adr_arg_type arg_type;
    int a;
    union
    {
        int b;
        T x;
    };
};

template<class T> struct ADR_TAPE;

template<class T>
struct ADR_T
{
    T x=0;
    ADR_TAPE<T>* tape=0;
    int i=0;

    ADR_T()=default;
    ADR_T(T t):x(t){}
    ADR_T(const ADR_T&)=default;
    ADR_T(T t,ADR_TAPE<T>* p,int j):x(t),tape(p),i(j) {}
    explicit operator bool() const {return x!=0;}
    explicit operator int() const {return x;}
};
template<class T>
struct ADR_TAPE
{
    ARRAY<T> x,dx;
    ARRAY<ADR_OP<T> > op;

    ADR_TAPE()
        :x(1),op(1,init_new,{op_nop,adr_arg_type::arg_none})
    {
    }

    void Reset()
    {
        x.Remove_All();
        dx.Remove_All();
        op.Remove_All();
    }

    // Compute the tape with respect to z.
    void Compute_Value()
    {
        for(int i=0;i<op.m;i++)
        {
            ADR_OP w=op(i);
            if(w.op_type==op_nop) continue;
            T a=0,b=0;
            Fill_Args(a,b,w);
            x(i)=TABLE<T>::entries[w.op_type].op(a,b);
        }
    }

    // Compute gradient with respect to z.
    void Compute_Gradient(ADR_T<T> z)
    {
        dx.Resize(x.m,init_all,0);
        dx(z.i)=1;
        for(int i=z.i;i>0;i--)
        {
            ADR_OP w=op(i);
            T a=0,*da=0,b=0,*db=0;
            Fill_Args(a,da,b,db,w);
            TABLE<T>::entries[w.op_type].rev_op(a,*da,b,*db,x(i),dx(i));
        }
    }

    void Fill_Args(T& a,T& b,ADR_OP<T> op)
    {
        switch(op.arg_type)
        {
            case adr_arg_type::arg_none:break;
            case adr_arg_type::arg_a:a=x(op.a);break;
            case adr_arg_type::arg_x:a=op.x;break;
            case adr_arg_type::arg_ab:a=x(op.a);b=x(op.b);break;
            case adr_arg_type::arg_ax:a=x(op.a);b=op.x;break;
            case adr_arg_type::arg_xa:a=op.x;b=x(op.a);break;
        }
    }

    void Fill_Args(T& a,T*& da,T& b,T*& db,ADR_OP<T> op)
    {
        switch(op.arg_type)
        {
            case adr_arg_type::arg_none:break;
            case adr_arg_type::arg_a:a=x(op.a);da=&dx(op.a);break;
            case adr_arg_type::arg_x:a=op.x;da=&dx(0);break;
            case adr_arg_type::arg_ab:a=x(op.a);da=&dx(op.a);b=x(op.b);db=&dx(op.b);break;
            case adr_arg_type::arg_ax:a=x(op.a);da=&dx(op.a);b=op.x;db=&dx(0);break;
            case adr_arg_type::arg_xa:a=op.x;da=&dx(0);b=x(op.a);db=&dx(op.a);break;
        }
    }

    ADR_T<T> Op(adr_op_type tp,T v,ADR_T<T> a,ADR_T<T> b)
    {
        if(!a.tape) return Op(tp,v,a.x,b);
        if(!b.tape) return Op(tp,v,a,b.x);
        int r=x.Append(v);
        op.Append({tp,adr_arg_type::arg_ab,a.i,{.b=b.i}});
        return {v,this,r};
    }

    ADR_T<T> Op(adr_op_type tp,T v,ADR_T<T> a,T b)
    {
        if(!a.tape) return v;
        int r=x.Append(v);
        op.Append({tp,adr_arg_type::arg_ax,a.i,{.x=b}});
        return {v,this,r};
    }

    ADR_T<T> Op(adr_op_type tp,T v,T a,ADR_T<T> b)
    {
        if(!b.tape) return v;
        int r=x.Append(v);
        op.Append({tp,adr_arg_type::arg_xa,b.i,{.x=a}});
        return {v,this,r};
    }

    ADR_T<T> Op(adr_op_type tp,T v,T a,T b)
    {
        return v;
    }

    ADR_T<T> Op(adr_op_type tp,T v,ADR_T<T> a)
    {
        if(!a.tape) return v;
        int r=x.Append(v);
        op.Append({tp,adr_arg_type::arg_a,a.i});
        return {v,this,r};
    }

    ADR_T<T> Op(adr_op_type tp,T v,T a)
    {
        return v;
    }

    ADR_T<T> Var(T a)
    {
        int r=x.Append(a);
        op.Append({op_nop,ADR::adr_arg_type::arg_none});
        return {a,this,r};
    }

    template<int d>
    VECTOR<ADR_T<T>,d> Var(const VECTOR<T,d>& a)
    {
        VECTOR<ADR_T<T>,d> u;
        for(int i=0;i<d;i++)
            u(i)=Var(a(i));
        return u;
    }

    template<int m,int n>
    MATRIX<ADR_T<T>,m,n> Var(const MATRIX<T,m,n>& a)
    {
        MATRIX<ADR_T<T>,m,n> r;
        for(int i=0;i<m;i++)
            for(int j=0;j<n;j++)
                r(i,j)=Var(a(i,j));
        return r;
    }

    void Set_Var(ADR_T<T>& z,T a)
    {
        z.x=a;
        x(z.i)=a;
    }

    template<int d>
    void Set_Var(VECTOR<ADR_T<T>,d>& z,const VECTOR<T,d>& a)
    {
        for(int i=0;i<d;i++)
            Set_Var(z(i),a(i));
    }

    template<int m,int n>
    void Set_Var(MATRIX<ADR_T<T>,m,n>& z,const MATRIX<T,m,n>& a)
    {
        for(int i=0;i<m;i++)
            for(int j=0;j<n;j++)
                Set_Var(z(i,j),a(i,j));
    }

    ADR_T<T> Const(T a)
    {
        return ADR_T<T>(a);
    }

    template<int d>
    VECTOR<ADR_T<T>,d> Const(const VECTOR<T,d>& a)
    {
        return VECTOR<ADR_T<T>,d>(a);
    }

    template<int m,int n>
    MATRIX<ADR_T<T>,m,n> Const(const MATRIX<T,m,n>& a)
    {
        return MATRIX<ADR_T<T>,m,n>(a);
    }

    void Print() const
    {
        for(int i=0;i<op.m;i++)
        {
            ADR_OP w=op(i);
            if(w.op_type==op_nop) continue;
            printf("x[%i] = ",i);
            char arg0[24],arg1[24];
            switch(w.arg_type)
            {
                case adr_arg_type::arg_none:break;
                case adr_arg_type::arg_a:sprintf(arg0,"x[%i]",w.a);break;
                case adr_arg_type::arg_x:sprintf(arg0,"%.16g",w.x);break;
                case adr_arg_type::arg_ab:
                    sprintf(arg0,"x[%i]",w.a);
                    sprintf(arg1,"x[%i]",w.b);
                    break;
                case adr_arg_type::arg_ax:
                    sprintf(arg0,"x[%i]",w.a);
                    sprintf(arg1,"%.16g",w.x);
                    break;
                case adr_arg_type::arg_xa:
                    sprintf(arg0,"%.16g",w.x);
                    sprintf(arg1,"x[%i]",w.a);
                    break;
            }
            printf(TABLE<T>::entries[w.op_type].pattern,arg0,arg1);
            printf(";\n");
        }
    }
};

template<class T>
T Value(ADR_T<T> a)
{
    if(!a.tape) return a.x;
    return a.tape->x(a.i);
}

template<class T,int d>
VECTOR<T,d> Value(const VECTOR<ADR_T<T>,d>& a)
{
    VECTOR<T,d> u;
    for(int i=0;i<d;i++)
        u(i)=Value(a(i));
    return u;
}

template<class T,int m,int n>
MATRIX<T,m,n> Value(const MATRIX<ADR_T<T>,m,n>& a)
{
    MATRIX<T,m,n> r;
    for(int i=0;i<m;i++)
        for(int j=0;j<n;j++)
            r(i,j)=Value(a(i,j));
    return r;
}

template<class T>
T Diff(ADR_T<T> a)
{
    return a.tape->dx(a.i);
}

template<class T,int d>
VECTOR<T,d> Diff(const VECTOR<ADR_T<T>,d>& a)
{
    VECTOR<T,d> u;
    for(int i=0;i<d;i++)
        u(i)=a(i).tape->dx(a(i).i);
    return u;
}

template<class T,int m,int n>
MATRIX<T,m,n> Diff(const MATRIX<ADR_T<T>,m,n>& a)
{
    MATRIX<T,m,n> r;
    for(int i=0;i<m;i++)
        for(int j=0;j<n;j++)
            r(i,j)=a(i,j).tape->dx(a(i,j).i);
    return r;
}

#define INST_OP1(op,type)                                               \
    template<class T> ADR_T<T> operator op(ADR_T<T> a)                  \
    {T v=op(a.x);return a.tape?a.tape->Op(type,v,a):ADR_T<T>{v};}

#define INST_HELP1(func,type)                                           \
    template<class T> ADR_T<T> func(ADR_T<T> a)                         \
    {T v=func(a.x);return a.tape?a.tape->Op(type,v,a):ADR_T<T>(v);}

#define INST_OP2(op,type)                                               \
    template<class T> ADR_T<T> operator op(ADR_T<T> a,T b)              \
    {                                                                   \
        T v=a.x op b;                                                   \
        if(!a.tape) return ADR_T<T>{v};                                 \
        return a.tape->Op(type,v,a,b);                                  \
    }                                                                   \
    template<class T> ADR_T<T> operator op(ADR_T<T> a,int b)            \
    {                                                                   \
        T v=a.x op b;                                                   \
        if(!a.tape) return ADR_T<T>{v};                                 \
        return a.tape->Op(type,v,a,b);                                  \
    }                                                                   \
    template<class T> ADR_T<T> operator op(T a,ADR_T<T> b)              \
    {                                                                   \
        T v=a op b.x;                                                   \
        if(!b.tape) return ADR_T<T>{v};                                 \
        return b.tape->Op(type,v,a,b);                                  \
    }                                                                   \
    template<class T> ADR_T<T> operator op(int a,ADR_T<T> b)            \
    {                                                                   \
        T v=a op b.x;                                                   \
        if(!b.tape) return ADR_T<T>{v};                                 \
        return b.tape->Op(type,v,a,b);                                  \
    }                                                                   \
    template<class T> ADR_T<T> operator op(ADR_T<T> a,ADR_T<T> b)       \
    {                                                                   \
        if(!a.tape) return a.x op b;                                    \
        return a.tape->Op(type,a.x op b.x,a,b);                         \
    }

#define INST_HELP2(func,type)                                   \
    template<class T> ADR_T<T> func(ADR_T<T> a,T b)             \
    {                                                           \
        T v=func(a.x,b);                                        \
        if(!a.tape) return ADR_T<T>{v};                         \
        return a.tape->Op(type,v,a,b);                          \
    }                                                           \
    template<class T> ADR_T<T> func(ADR_T<T> a,int b)           \
    {                                                           \
        T v=func(a.x,b);                                        \
        if(!a.tape) return ADR_T<T>{v};                         \
        return a.tape->Op(type,v,a,b);                          \
    }                                                           \
    template<class T> ADR_T<T> func(T a,ADR_T<T> b)             \
    {                                                           \
        T v=func(a,b.x);                                        \
        if(!b.tape) return ADR_T<T>{v};                         \
        return b.tape->Op(type,v,a,b);                          \
    }                                                           \
    template<class T> ADR_T<T> func(int a,ADR_T<T> b)           \
    {                                                           \
        T v=func(a,b.x);                                        \
        if(!b.tape) return ADR_T<T>{v};                         \
        return b.tape->Op(type,v,a,b);                          \
    }                                                           \
    template<class T> ADR_T<T> func(ADR_T<T> a,ADR_T<T> b)      \
    {                                                           \
        if(!a.tape) return func(a.x,b);                         \
        return a.tape->Op(type,func(a.x,b.x),a,b);              \
    }

INST_OP1(+,op_copy);
INST_OP1(-,op_neg);
INST_OP2(+,op_add);
INST_OP2(-,op_sub);
INST_OP2(*,op_mul);
INST_OP2(/,op_div);

template<class T> ADR_T<T>& operator+= (ADR_T<T>& a,ADR_T<T> b) {return a=a+b;}
template<class T> ADR_T<T>& operator-= (ADR_T<T>& a,ADR_T<T> b) {return a=a-b;}
template<class T> ADR_T<T>& operator*= (ADR_T<T>& a,ADR_T<T> b) {return a=a*b;}
template<class T> ADR_T<T>& operator/= (ADR_T<T>& a,ADR_T<T> b) {return a=a/b;}
template<class T> ADR_T<T>& operator+= (ADR_T<T>& a,T b) {return a=a+b;}
template<class T> ADR_T<T>& operator-= (ADR_T<T>& a,T b) {return a=a-b;}
template<class T> ADR_T<T>& operator*= (ADR_T<T>& a,T b) {return a=a*b;}
template<class T> ADR_T<T>& operator/= (ADR_T<T>& a,T b) {return a=a/b;}

INST_HELP1(Inverse,op_inv);
INST_HELP1(sqrt,op_sqrt);
INST_HELP1(exp,op_exp);
INST_HELP1(log,op_ln);
INST_HELP1(sin,op_sin);
INST_HELP1(cos,op_cos);
INST_HELP1(tan,op_tan);
INST_HELP1(asin,op_asin);
INST_HELP1(acos,op_acos);
INST_HELP1(atan,op_atan);
INST_HELP1(sinh,op_sinh);
INST_HELP1(cosh,op_cosh);
INST_HELP1(tanh,op_tanh);
INST_HELP1(asinh,op_asinh);
INST_HELP1(acosh,op_acosh);
INST_HELP1(atanh,op_atanh);
INST_HELP1(floor,op_floor);
INST_HELP1(ceil,op_ceil);
INST_HELP1(rint,op_rint);
INST_HELP2(atan2,op_atan2);
INST_HELP2(hypot,op_hypot);

#undef INST_HELP1
#undef INST_HELP2
}
template<class T> struct SCALAR_POLICY<ADR::ADR_T<T> > {typedef ADR::ADR_T<T> TYPE;};
}

#endif
