//#####################################################################
// Copyright 2018.  Ounan Ding, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
// Class Diff_Test
//#####################################################################
#ifndef __Diff_Test__
#define __Diff_Test__

#include <Core/Arrays/ARRAY.h>
#include <Core/Data_Structures/HASHTABLE.h>
#include <Core/Log/LOG.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/MATRIX_MXN.h>
#include <Core/Tensors/TENSORS_FORWARD.h>
#include <Core/Vectors/VECTOR.h>
#include <cfloat>
#include <cmath>
#include <string>
namespace PhysBAM{
using ::std::abs;

extern int diff_test_stage;

namespace DIFF_TEST_HELPERS{
struct D
{
    template<class T,class T_ARRAY>
    static auto Rem_Expr(const ARRAY_BASE<T,T_ARRAY>& a)
    {
        return VECTOR_COPY<T_ARRAY>(a.Derived());
    }

    template<class T>
    static enable_if_t<!IS_ARRAY<T>::value,const T&> Rem_Expr(const T& a){return a;}
    
    template<class A,class B> static auto
    Dot(const A& a,const B& b)
    {
        if constexpr(IS_SCALAR<B>::value)
            return a*b;
        else if constexpr(IS_ARRAY<B>::value)
        {
            if constexpr(IS_ARRAY<A>::value)
            {
                assert(a.Size()==b.Size());
                decltype(Dot(a(0),b(0))) x={};
                if(!a.Size()) return x;
                x=Dot(a(0),b(0))*0;
                for(int i=0;i<a.Size();i++) x+=Dot(a(i),b(i));
                return x;
            }
            else if constexpr(IS_MATRIX<A>::value)
            {
                assert(a.Columns()==b.Size());
                ARRAY<decltype(Dot(a(0,0),b(0)))> x(a.Rows());
                if(!a.Rows()||!a.Columns()) return x;
                for(int i=0;i<a.Rows();i++)
                {
                    x(i)=Dot(a(i,0),b(0))*0;
                    for(int j=0;j<a.Columns();j++)
                        x(i)+=Dot(a(i,j),b(j));
                }
                return x;
            }
            else
            {
                assert(A::p==b.Size());
                MATRIX<decltype(Dot(a(0,0,0),b(0))),A::m,A::n> x;
                for(int i=0;i<a.m;i++)
                    for(int j=0;j<a.n;j++)
                    {
                        x(i,j)=Dot(a(i,j,0),b(0))*0;
                        for(int k=0;k<a.p;k++)
                            x(i,j)+=Dot(a(i,j,k),b(k));
                    }
                return x;
            }
        }
        else if constexpr(IS_MATRIX<B>::value)
        {
            if constexpr(IS_MATRIX<A>::value)
            {
                assert(a.Rows()==b.Rows());
                assert(a.Columns()==b.Columns());
                decltype(Dot(a(0,0),b(0,0))) x={};
                if(!a.Rows()||!a.Columns()) return x;
                x=Dot(a(0,0),b(0,0))*0;
                for(int i=0;i<a.Rows();i++)
                    for(int j=0;j<a.Columns();j++)
                        x+=Dot(a(i,j),b(i,j));
                return x;
            }
            else
            {
                assert(A::n==b.Rows());
                assert(A::p==b.Columns());
                VECTOR<decltype(Dot(a(0,0,0),b(0,0))),A::m> x;
                for(int i=0;i<a.m;i++)
                {
                    x(i)=Dot(a(i,0,0),b(0,0))*0;
                    for(int j=0;j<a.n;j++)
                        for(int k=0;k<a.p;k++)
                            x(i)+=Dot(a(i,j,k),b(j,k));
                }
                return x;
            }
        }
        else
        {
            static_assert((int)a.m==(int)b.m && (int)a.n==(int)b.n && (int)a.p==(int)b.p);
            decltype(Dot(a(0,0,0),b(0,0,0))) x={};
            if(!a.Rows()||!a.Columns()) return x;
            x=Dot(a(0,0,0),b(0,0,0))*0;
            for(int i=0;i<a.m;i++)
                for(int j=0;j<a.n;j++)
                    for(int k=0;k<a.p;k++)
                        x+=Dot(a(i,j,k),b(i,j,k));
            return x;
        }
    }

    template<size_t... n,class... T0,class... T1>
    static auto Dot_(std::index_sequence<n...>,const std::tuple<T0...>& a,const std::tuple<T1...>& b)
    {
        return (Dot(std::get<n>(a),std::get<n>(b))+...);
    }

    template<class... T0,class... T1>
    static auto Dot(const std::tuple<T0...>& a,const std::tuple<T1...>& b)
    {
        return Dot_(std::make_index_sequence<sizeof...(T0)>(),a,b);
    }

    template<class T>
    static auto Magnitude(const T& a)
    {
        return sqrt(Dot(a,a));
    }

    template<size_t... n,class... T0,class... T1,class... T2,class... T3>
    static auto Test_(std::index_sequence<n...>,
        const std::tuple<T0...>& a,const std::tuple<T1...>& da,
        const std::tuple<T2...>& b,const std::tuple<T3...>& db)
    {
        auto d=std::make_tuple((Rem_Expr(std::get<n>(a)-std::get<n>(b)))...);
        auto s=std::make_tuple((Rem_Expr(std::get<n>(da)+std::get<n>(db)))...);
        auto exact=Rem_Expr(D::Dot(s,d)/2);
        auto eps=D::Magnitude(d);
        return std::make_pair(exact,eps);
    }

    template<class... T0,class... T1,class... T2,class... T3>
    static auto Test(const std::tuple<T0...>& a,const std::tuple<T1...>& da,
        const std::tuple<T2...>& b,const std::tuple<T3...>& db)
    {
        return Test_(std::make_index_sequence<sizeof...(T0)>(),a,da,b,db);
    }

    template<size_t... n,class... T0,class T1,class... T2,class T3>
    static auto Mul_Test_(std::index_sequence<n...>,
        const std::tuple<T0...>& a,const T1& da,
        const std::tuple<T2...>& b,const T3& db)
    {
        auto d=std::make_tuple((Rem_Expr(std::get<n>(a)-std::get<n>(b)))...);
        auto exact=Rem_Expr((da+db)/2);
        auto eps=D::Magnitude(d);
        return std::make_pair(exact,eps);
    }

    template<class... T0,class T1,class... T2,class T3>
    static auto Mul_Test(const std::tuple<T0...>& a,const T1& da,
        const std::tuple<T2...>& b,const T3& db)
    {
        return Mul_Test_(std::make_index_sequence<sizeof...(T0)>(),a,da,b,db);
    }
    
    template<size_t... n,class T,class ...Args>
    static auto Pack_(std::index_sequence<n...>,const T& Y,const Args&... args)
    {
        auto t=std::make_tuple(args...);
        return std::make_tuple(Y,
            std::make_tuple(std::get<n*2>(t)...),
            std::make_tuple(std::get<n*2+1>(t)...));
    }

    template<class T,class ...Args>
    static auto Pack(const T& Y,const Args&... args)
    {
        return Pack_(std::make_index_sequence<sizeof...(Args)/2>(),Y,args...);
    }
    
    template<int half,size_t... n,class ...Args>
    static auto Pack6_(std::index_sequence<n...>,const Args&... args)
    {
        auto t=std::make_tuple(args...);
        return std::make_tuple(std::get<0>(t),
            std::make_tuple(std::get<n*2+1>(t)...),
            std::make_tuple(std::get<n*2+2>(t)...),
            std::get<half>(t),
            std::make_tuple(std::get<half+n*2+1>(t)...),
            std::make_tuple(std::get<half+n*2+2>(t)...));
    }

    template<class ...Args>
    static auto Pack6(const Args&... args)
    {
        static_assert(sizeof...(Args)%4==2);
        return Pack6_<sizeof...(Args)/2>(
            std::make_index_sequence<sizeof...(Args)/4>(),args...);
    }

    template<int half,size_t... n,class ...Args>
    static auto Pack6m_(std::index_sequence<n...>,const Args&... args)
    {
        auto t=std::make_tuple(args...);
        return std::make_tuple(std::get<0>(t),std::get<1>(t),
            std::make_tuple(std::get<n+2>(t)...),
            std::get<half>(t),std::get<half+1>(t),
            std::make_tuple(std::get<half+n+2>(t)...));
    }

    template<class ...Args>
    static auto Pack6m(const Args&... args)
    {
        return Pack6m_<sizeof...(Args)/2>(
            std::make_index_sequence<sizeof...(Args)/2-2>(),args...);
    }
};

template<class PACK>
HASHTABLE<std::string,PACK>& Diff_Test_Store()
{
    static HASHTABLE<std::string,PACK> h;
    return h;
}

template<class F,class X,class T>
void Report_Error(std::ostream& os,const char* name,const F& diff,const X& exact,T eps)
{
    T d=D::Magnitude(diff)/eps;
    T e=D::Magnitude(exact)/eps;
    T r=D::Magnitude(diff-exact)/eps;
    LOG::fprintf(os,"diff test %s:  diff %.16P  exact %.16P   abs-error %P   rel-error %P\n",
        name,d,e,r,r/max(max(d,e),(T)FLT_MIN));
}

// Diff_Test("f",0/1,f,x,df_dx,y,df_dy,...)
template<class T,class ...Args> 
void Diff_Test(std::ostream& os,int step,const char* name,const T& Y,const Args&... args)
{
    if(step!=0 && step!=1) return;
    auto cur=D::Pack(Y,args...);
    typedef decltype(cur) PACK;
    auto& hash=Diff_Test_Store<PACK>();

    if(step==0) hash.Set(name,cur);
    else if(step==1)
    {
        const auto& old=hash.Get(name);
        auto diff=std::get<0>(cur)-std::get<0>(old);
        auto [exact,eps]=D::Test(std::get<1>(cur),std::get<2>(cur),std::get<1>(old),std::get<2>(old));
        Report_Error(os,name,diff,exact,eps);
        hash.Delete(name);
    }
}
template<class T,class ...Args> 
void Diff_Test(int step,const char* name,const T& Y,const Args&... args)
{Diff_Test(LOG::cout,step,name,Y,args...);}

// Diff_Test("f",f0,x0,df_dx0,y0,df_dy0,...,f1,x1,df_dx1,y1,df_dy1,...)
template<class ...Args> 
void Diff_Test(std::ostream& os,const char* name,const Args&... args)
{
    auto [y0,x0,dy0,y1,x1,dy1]=D::Pack6(args...);
    auto diff=y1-y0;
    auto [exact,eps]=D::Test(x1,dy1,x0,dy0);
    Report_Error(os,name,diff,exact,eps);
};
template<class ...Args> 
void Diff_Test(const char* name,const Args&... args)
{Diff_Test(LOG::cout,name,args...);}

// Diff_Times_Test("f",0/1,f,df_dxi_times_dxi,dx0,dx1,...)
template<class T,class ...Args> 
void Diff_Times_Test(std::ostream& os,int step,const char* name,const T& Y,const T& dY,const Args&... args)
{
    if(step!=0 && step!=1) return;
    auto cur=std::make_tuple(Y,dY);
    typedef decltype(cur) PACK;
    auto& hash=Diff_Test_Store<PACK>();

    if(step==0) hash.Set(name,cur);
    else if(step==1)
    {
        const auto& old=hash.Get(name);
        auto diff=std::get<0>(cur)-std::get<0>(old);
        auto exact=D::Rem_Expr((std::get<1>(cur)+std::get<1>(old))/2);
        auto eps=sqrt((D::Dot(args,args)+...));
        Report_Error(os,name,diff,exact,eps);
        hash.Delete(name);
    }
}
template<class T,class ...Args> 
void Diff_Times_Test(int step,const char* name,const T& Y,const T& dY,const Args&... args)
{Diff_Times_Test(LOG::cout,step,name,Y,dY,args...);}

// Diff_Test("f",f0,df_dx0_times_dx0,x0a,x0b,...,f1,df_dx1_times_dx1,x1a,x1b,...)
template<class ...Args> 
void Diff_Times_Test(std::ostream& os,const char* name,const Args&... args)
{
    auto [y0,dy0,x0,y1,dy1,x1]=D::Pack6m(args...);
    auto diff=y1-y0;
    auto [exact,eps]=D::Mul_Test(x1,dy1,x0,dy0);
    Report_Error(os,name,diff,exact,eps);
};
template<class ...Args> 
void Diff_Times_Test(const char* name,const Args&... args)
{Diff_Times_Test(LOG::cout,name,args...);}
}
using DIFF_TEST_HELPERS::Diff_Test;
using DIFF_TEST_HELPERS::Diff_Times_Test;
}
#endif
