//#####################################################################
// Copyright 2012.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYREGHT.txt.
//#####################################################################
// Class PROGRAM_SCALAR
//#####################################################################
#ifndef __PROGRAM_SCALAR__
#define __PROGRAM_SCALAR__

#include <Core/Math_Tools/cube.h>
#include <Core/Math_Tools/sqr.h>
#include <Core/Symbolics/INSTRUCTION.h>
#include <Core/Vectors/SCALAR_POLICY.h>
#include <cmath>
#include <functional>
#include <iosfwd>

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{

template<class T> struct PROGRAM_COMPILER;

namespace NS_PROGRAM_SCALAR
{
using ::std::sin;
using ::std::cos;
using ::std::tan;
using ::std::asin;
using ::std::acos;
using ::std::atan;
using ::std::atan2;
using ::std::exp;
using ::std::log;
using ::std::abs;
using ::std::pow;
using ::std::sqrt;
using ::PhysBAM::sqr;
using ::PhysBAM::cube;
using ::std::hypot;
using ::std::floor;
using ::std::ceil;
using ::std::rint;

template<class T>
class PROGRAM_SCALAR
{
public:
    typedef T RAW_SCALAR;
    PROGRAM_COMPILER<T>* pc=nullptr;
    REG var;

    static PROGRAM_SCALAR<T> Number(PROGRAM_COMPILER<T>& pc,T a);
    void Print() const;
};

template<class T>
PROGRAM_SCALAR<T> Make_Scalar(PROGRAM_COMPILER<T>& pc,const PROGRAM_SCALAR<T>& a){return a;}

template<class T>
PROGRAM_SCALAR<T> Make_Scalar(PROGRAM_COMPILER<T>& pc,const ND<T>& a)
{return PROGRAM_SCALAR<T>::Number(pc,a);}

template<class T>
PROGRAM_SCALAR<T> Make_Output(const PROGRAM_SCALAR<T>& a)
{return {a.pc,a.pc->Make_Output(a.var)};}

template<class T,int d>
VECTOR<PROGRAM_SCALAR<T>,d> Make_Output(const VECTOR<PROGRAM_SCALAR<T>,d>& a)
{
    VECTOR<PROGRAM_SCALAR<T>,d> r(a);
    for(auto&i:r) i=Make_Output(i);
    return r;
}

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

#define MK1(f) \
    template<class T> PROGRAM_SCALAR<T> f(const PROGRAM_SCALAR<T>& a);

#define MK2(f) \
    template<class T> PROGRAM_SCALAR<T> f(const PROGRAM_SCALAR<T>& a,const PROGRAM_SCALAR<T>& b); \
    template<class T> PROGRAM_SCALAR<T> f(const PROGRAM_SCALAR<T>& a,const ND<T>& b) \
    {return f(PROGRAM_SCALAR<T>(a),Make_Scalar(a.pc,b));}               \
    template<class T> PROGRAM_SCALAR<T> f(const ND<T>& a,const PROGRAM_SCALAR<T>& b) \
    {return f(Make_Scalar(b.pc,a),PROGRAM_SCALAR<T>(b));}

#define MK3(f)                                                          \
    template<class T> PROGRAM_SCALAR<T> f(const PROGRAM_SCALAR<T>& a,const PROGRAM_SCALAR<T>& b,const PROGRAM_SCALAR<T>& c); \
    template<class T,class B,class C> PROGRAM_SCALAR<T> f(const PROGRAM_SCALAR<T>& a,const B& b,const C& c) \
    {return f(a,Make_Scalar(a.pc,b),Make_Scalar(a.pc,c));} \
    template<class T,class C> PROGRAM_SCALAR<T> f(const ND<T>& a,const PROGRAM_SCALAR<T>& b,const C& c) \
    {return f(Make_Scalar(b.pc,a),b,Make_Scalar(b.pc,c));} \
    template<class T> PROGRAM_SCALAR<T> f(const ND<T>& a,const ND<T>& b,const PROGRAM_SCALAR<T>& c) \
    {return f(Make_Scalar(c.pc,a),Make_Scalar(c.pc,b),c);} \

MK1(operator+);
MK1(operator-);
MK2(operator+);
MK2(operator-);
MK2(operator*);
MK2(operator/);

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

MK2(operator<);
MK2(operator>);
MK2(operator<=);
MK2(operator>=);
MK2(operator==);
MK2(operator!=);
MK3(cond);

template<class T> PROGRAM_SCALAR<T>
cond_lazy(const PROGRAM_SCALAR<T>& a,
    std::function<PROGRAM_SCALAR<ND<T>>()> b,
    std::function<PROGRAM_SCALAR<ND<T>>()> c);
template<class T> ARRAY<PROGRAM_SCALAR<T>>
cond_lazy(const PROGRAM_SCALAR<T>& a,
    std::function<ARRAY<PROGRAM_SCALAR<ND<T>>>()> b,
    std::function<ARRAY<PROGRAM_SCALAR<ND<T>>>()> c);

MK1(abs);
MK1(sqrt);
MK1(exp);
MK1(log);
MK1(sin);
MK1(cos);
MK1(tan);
MK1(asin);
MK1(acos);
MK1(atan);
MK1(floor);
MK1(ceil);
MK1(rint);
MK1(Inverse);
MK2(atan2);
MK2(hypot);
MK2(pow);
MK2(min);
MK2(max);

template<class T> std::ostream& operator<<(std::ostream& o, const PROGRAM_SCALAR<T>& s);
template<class T> std::istream& operator>>(std::istream& i, PROGRAM_SCALAR<T>& s);

MK1(Magnitude_Squared);
MK2(Dot_Product);
MK3(wrap);

template<class T> PROGRAM_SCALAR<T>
Diff(const PROGRAM_SCALAR<T>& a,MEM_IN var);
template<class T> void
Diff(PROGRAM_SCALAR<T>* r,const PROGRAM_SCALAR<T>* a,int m,MEM_IN var);
template<class T> void
Diff(PROGRAM_SCALAR<T>* r,const PROGRAM_SCALAR<T>& a,const MEM_IN* var,int n);
template<class T> void
Diff(PROGRAM_SCALAR<T>* r,int si,int sj,const PROGRAM_SCALAR<T>* a,int m,const MEM_IN* var,int n);

template<class T,int d> VECTOR<PROGRAM_SCALAR<T>,d>
Diff(const VECTOR<PROGRAM_SCALAR<T>,d>& a,MEM_IN var)
{
    VECTOR<PROGRAM_SCALAR<T>,d> r;
    Diff(&r(0),&a(0),d,var);
    return r;
}

template<class T,int m,int n> MATRIX<PROGRAM_SCALAR<T>,m,n>
Diff(const MATRIX<PROGRAM_SCALAR<T>,m,n>& a,MEM_IN var)
{
    MATRIX<PROGRAM_SCALAR<T>,m,n> r;
    Diff(&r(0,0),&a(0,0),m*n,var);
    return r;
}

template<class T,int m> SYMMETRIC_MATRIX<PROGRAM_SCALAR<T>,m>
Diff(const SYMMETRIC_MATRIX<PROGRAM_SCALAR<T>,m>& a,MEM_IN var)
{
    SYMMETRIC_MATRIX<PROGRAM_SCALAR<T>,m> r;
    Diff(&r(0,0),&a(0,0),m*(m+1)/2,var);
    return r;
}

template<class T,int d> VECTOR<PROGRAM_SCALAR<T>,d>
Gradient(const PROGRAM_SCALAR<T>& a,const VECTOR<MEM_IN,d>& vars)
{
    VECTOR<PROGRAM_SCALAR<T>,d> r;
    Diff(&r(0),a,&vars(0),d);
    return r;
}

template<class T,int m,int n> MATRIX<PROGRAM_SCALAR<T>,m,n>
Gradient(const VECTOR<PROGRAM_SCALAR<T>,m>& u,const VECTOR<MEM_IN,n>& vars)
{
    MATRIX<PROGRAM_SCALAR<T>,m,n> r;
    assert(&r(1,0)-&r(0,0)==1 && &r(0,1)-&r(0,0)==m);
    Diff(&r(0,0),1,m,&u(0),m,&vars(0),n);
    return r;
}

template<class T,int d> PROGRAM_SCALAR<T>
Divergence(const VECTOR<PROGRAM_SCALAR<T>,d>& u,const VECTOR<MEM_IN,d>& vars)
{
    PROGRAM_SCALAR<T> r=Diff(u(0),vars(0));
    for(int i=1;i<d;i++) r+=Diff(u(i),vars(i));
    return r;
}

// Sum on first index
template<class T,int m,int n> VECTOR<PROGRAM_SCALAR<T>,n>
Divergence_L(const MATRIX<PROGRAM_SCALAR<T>,m,n>& u,const VECTOR<MEM_IN,m>& vars)
{
    VECTOR<PROGRAM_SCALAR<T>,n> r;
    for(int i=0;i<n;i++) r(i)=Divergence(u.Column(i),vars);
    return r;
}

// Sum on second index
template<class T,int m,int n> VECTOR<PROGRAM_SCALAR<T>,m>
Divergence_R(const MATRIX<PROGRAM_SCALAR<T>,m,n>& u,const VECTOR<MEM_IN,n>& vars)
{
    VECTOR<PROGRAM_SCALAR<T>,n> r;
    for(int i=0;i<n;i++) r(i)=Divergence(u.Row(i),vars);
    return r;
}

template<class T,int d> PROGRAM_SCALAR<T>
Laplacian(const PROGRAM_SCALAR<T>& a,const VECTOR<MEM_IN,d>& vars)
{
    PROGRAM_SCALAR<T> r=Diff(Diff(a,vars(0)),vars(0));
    for(int i=1;i<d;i++) r+=Diff(Diff(a,vars(i)),vars(i));
    return r;
}

template<class T,int m,int n> VECTOR<PROGRAM_SCALAR<T>,m>
Laplacian(const VECTOR<PROGRAM_SCALAR<T>,m>& u,const VECTOR<MEM_IN,n>& vars)
{
    VECTOR<PROGRAM_SCALAR<T>,m> r;
    for(int i=0;i<m;i++) r(i)=Laplacian(u(i),vars);
    return r;
}
}
using NS_PROGRAM_SCALAR::PROGRAM_SCALAR;

template<class T> struct SCALAR_POLICY<PROGRAM_SCALAR<T>> {typedef PROGRAM_SCALAR<T> TYPE;};

#undef MK1
#undef MK2
#undef MK3
}
#endif
