//#####################################################################
// Copyright 2013.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
// Class AUTO_DIFF_REV
//##################################################################### 
#ifndef __AUTO_DIFF_REV__
#define __AUTO_DIFF_REV__
#include <Core/Vectors/VECTOR.h>

namespace PhysBAM{


template<class T,class AD>
struct AUTO_DIFF_BASE
{
    enum ADR {adr = 1};
};

template<class T>
struct AUTO_DIFF_VAR: AUTO_DIFF_BASE<T,AUTO_DIFF_VAR<T> >
{
    typedef T TYPE;
    T x;
    T dx=T();

    AUTO_DIFF_VAR(const T& t)
        :x(t)
    {
    }

    void Update_Diff(){}
};

template<class T> AUTO_DIFF_VAR<T> Make_ADV(const T& a){return {a};}

template<class U,class OP>
struct AUTO_DIFF_REV_1: AUTO_DIFF_BASE<decltype(OP::f(typename std::remove_reference_t<U>::TYPE())),AUTO_DIFF_REV_1<U,OP> >
{
    U u;
    typedef decltype(OP::f(u.x)) T;
    typedef T TYPE;
    T x,dx=T();
    bool updated=false;

    AUTO_DIFF_REV_1(U& uu)
        :u(uu),x(OP::f(u.x))
    {
    }

    void Update_Diff()
    {
        if(!updated)
        {
            OP::df_t(u.dx,dx,u.x,x);
            u.Update_Diff();
            updated=true;
        }
    }

    void Diff()
    {
        dx+=1;
        Update_Diff();
    }
};

template<class U,class V,class OP>
struct AUTO_DIFF_REV_2:
        AUTO_DIFF_BASE<decltype(
            OP::f(typename std::remove_reference_t<U>::TYPE(),
                typename std::remove_reference_t<V>::TYPE())),
            AUTO_DIFF_REV_2<U,V,OP> >
{
    U u;
    V v;
    typedef decltype(OP::f(u.x,v.x)) T;
    typedef T TYPE;
    T x,dx=T();
    bool updated=false;

    AUTO_DIFF_REV_2(U& uu,V& vv)
        :u(uu),v(vv),x(OP::f(u.x,v.x))
    {
    }

    void Update_Diff()
    {
        if(!updated)
        {
            OP::df_t(u.dx,v.dx,dx,u.x,v.x,x);
            u.Update_Diff();
            v.Update_Diff();
            updated=true;
        }
    }

    void Diff()
    {
        dx+=1;
        Update_Diff();
    }
};

#define ADR_DEF_1_(func,op,A,B) template<class AD0> std::enable_if_t<AD0::adr,AUTO_DIFF_REV_1<AD0 A,op> > func (AD0 B a) {return {a};}
#define ADR_DEF_1(func,op) ADR_DEF_1_(func,op,&,&);ADR_DEF_1_(func,op,,&&);
#define ADR_DEF_2_(func,op,A,B,C,D) template<class AD0,class AD1> std::enable_if_t<AD0::adr&&AD1::adr,AUTO_DIFF_REV_2<AD0 A,AD1 C,op> > func (AD0 B a, AD1 D b) {return {a,b};}
#define ADR_DEF_2(func,op) ADR_DEF_2_(func,op,&,&,&,&);ADR_DEF_2_(func,op,,&&,&,&);ADR_DEF_2_(func,op,&,&,,&&);ADR_DEF_2_(func,op,,&&,,&&);

#define ADR_DEF_M0_(func,op,A,B) AUTO_DIFF_REV_1<AD A,op> func () B {return {static_cast<AD&>(*this)};}
#define ADR_DEF_M0(func,op) ADR_DEF_M0_(func,op,&,&);ADR_DEF_M0_(func,op,,&&);
#define ADR_DEF_M1_(func,op,A,B,C,D) template<class AD0> auto func (AD0 D a) B -> std::enable_if_t<AD0::adr,AUTO_DIFF_REV_2<AD A,AD0 C,op> > {return {static_cast<AD&>(*this),a};}
#define ADR_DEF_M1(func,op) ADR_DEF_M1_(func,op,&,&,&,&);ADR_DEF_M1_(func,op,,&&,&,&);ADR_DEF_M1_(func,op,&,&,,&&);ADR_DEF_M1_(func,op,,&&,,&&);

struct ADR_NEG
{
    template<class T> static T f(T a){return -a;}
    template<class T> static void df_t(T& da,T dz,T a,T z){da-=dz;}
};
ADR_DEF_1(operator-,ADR_NEG);

struct ADR_ADD
{
    template<class T> static T f(T a,T b){return a+b;}
    template<class T> static void df_t(T& da,T& db,T dz,T a,T b,T z){da+=dz;db+=dz;}
};
ADR_DEF_2(operator+,ADR_ADD);

struct ADR_SUB
{
    template<class T> static T f(T a,T b){return a-b;}
    template<class T> static void df_t(T& da,T& db,T dz,T a,T b,T z){da+=dz;db-=dz;}
};
ADR_DEF_2(operator-,ADR_SUB);

struct ADR_MUL
{
    template<class T> static T f(T a,T b){return a*b;}
    template<class T> static void df_t(T& da,T& db,T dz,T a,T b,T z){da+=b*dz;db+=a*dz;}

    template<class T,int d> static VECTOR<T,d> f(VECTOR<T,d> a,T b){return a*b;}
    template<class T,int d> static void df_t(VECTOR<T,d>& da,T& db,VECTOR<T,d> dz,VECTOR<T,d> a,T b,VECTOR<T,d> z){da+=b*dz;db+=a.Dot(dz);}

    template<class T,int d> static VECTOR<T,d> f(T a,VECTOR<T,d> b){return a*b;}
    template<class T,int d> static void df_t(T& da,VECTOR<T,d>& db,VECTOR<T,d> dz,T a,VECTOR<T,d> b,VECTOR<T,d> z){da+=b.Dot(dz);db+=a*dz;}
};
ADR_DEF_2(operator*,ADR_MUL);

struct ADR_DIV
{
    template<class T> static T f(const T& a,const T& b){return a/b;}
    template<class T> static void df_t(T& da,T& db,T dz,T a,T b,T z){da+=dz/b;db-=z/b*dz;}

    template<class T,int d> static VECTOR<T,d> f(const VECTOR<T,d>& a,const T& b){return a/b;}
    template<class T,int d> static void df_t(VECTOR<T,d>& da,T& db,VECTOR<T,d> dz,VECTOR<T,d> a,T b,VECTOR<T,d> z){da+=dz/b;db-=z.Dot(dz)/b;}
};
ADR_DEF_2(operator/,ADR_DIV);

struct ADR_SIN
{
    template<class T> static T f(T a){return sin(a);}
    template<class T> static void df_t(T& da,T dz,T a,T z){da+=cos(a)*dz;}
};
ADR_DEF_1(sin,ADR_SIN);

struct ADR_COS
{
    template<class T> static T f(T a){return cos(a);}
    template<class T> static void df_t(T& da,T dz,T a,T z){da-=sin(a)*dz;}
};
ADR_DEF_1(cos,ADR_COS);

struct ADR_TAN
{
    template<class T> static T f(T a){return tan(a);}
    template<class T> static void df_t(T& da,T dz,T a,T z){da+=(z*z+1)*dz;}
};
ADR_DEF_1(tan,ADR_TAN);

struct ADR_EXP
{
    template<class T> static T f(T a){return exp(a);}
    template<class T> static void df_t(T& da,T dz,T a,T z){da+=z*dz;}
};
ADR_DEF_1(exp,ADR_EXP);

struct ADR_LOG
{
    template<class T> static T f(T a){return log(a);}
    template<class T> static void df_t(T& da,T dz,T a,T z){da+=dz/a;}
};
ADR_DEF_1(log,ADR_LOG);

struct ADR_DOT
{
    template<class TV> static auto f(TV a,TV b){return a.Dot(b);}
    template<class T,class TV> static void df_t(TV& da,TV& db,T dz,TV a,TV b,T z){da+=b*dz;db+=a*dz;}
};

struct ADR_MAG2
{
    template<class TV> static auto f(TV a){return a.Magnitude_Squared();}
    template<class T,class TV> static void df_t(TV& da,T dz,TV a,T z){da+=2*dz*a;}
};

struct ADR_MAG
{
    template<class TV> static auto f(TV a){return a.Magnitude();}
    template<class T,class TV> static void df_t(TV& da,T dz,TV a,T z){da+=dz/z;}
};

template<class T,int d,class AD>
struct AUTO_DIFF_BASE<VECTOR<T,d>,AD>
{
    enum ADR {adr = 1};

    ADR_DEF_M1(Dot,ADR_DOT);
    ADR_DEF_M0(Magnitude,ADR_MAG);
    ADR_DEF_M0(Magnitude_Squared,ADR_MAG2);
};

#undef ADR_DEF_1_
#undef ADR_DEF_1
#undef ADR_DEF_2_
#undef ADR_DEF_2
#undef ADR_DEF_M0_
#undef ADR_DEF_M0
#undef ADR_DEF_M1_
#undef ADR_DEF_M1

}
#endif
