//#####################################################################
// Copyright 2024, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Matrices/Eigenvalues.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>

namespace PhysBAM{

template<class T,int d>
DIAGONAL_MATRIX<T,d>
Eigenvalues(const SYMMETRIC_MATRIX<T,d>& A)
{
    if constexpr(d==0) return {};
    else if constexpr(d==1) return {A.x[0]};
    else if constexpr(d==2)
    {
        T da=0;
        if(A(1,0)!=0)
        {
            T theta=(T).5*(A(1,1)-A(0,0))/A(1,0),t=1/(abs(theta)+sqrt(1+sqr(theta)));
            if(theta<0) t=-t;
            da=t*A(1,0);
        }
        return DIAGONAL_MATRIX<T,2>(VECTOR<T,2>(A(0,0)-da,A(1,1)+da).Sorted().Reversed());
    }
    else if constexpr(d==3)
    {
        if constexpr(is_same<T,float>::value)
            return DIAGONAL_MATRIX<T,3>(SYMMETRIC_MATRIX<double,3>(A).Eigenvalues());
        else
        {
            // now T is double
            T m=((T)1/3)*(A(0,0)+A(1,1)+A(2,2));
            T a00=A(0,0)-m,a11=A(1,1)-m,a22=A(2,2)-m,a12_sqr=A(1,0)*A(1,0),a13_sqr=A(2,0)*A(2,0),a23_sqr=A(2,1)*A(2,1);
            T p=((T)1/6)*(a00*a00+a11*a11+a22*a22+2*(a12_sqr+a13_sqr+a23_sqr));
            T q=(T).5*(a00*(a11*a22-a23_sqr)-a11*a13_sqr-a22*a12_sqr)+A(1,0)*A(2,0)*A(2,1);
            T sqrt_p=sqrt(p),disc=p*p*p-q*q;
            T phi=((T)1/3)*atan2(sqrt(max((T)0,disc)),q),c=cos(phi),s=sin(phi);
            T sqrt_p_cos=sqrt_p*c,root_three_sqrt_p_sin=sqrt((T)3)*sqrt_p*s;
            VECTOR<T,3> lambda(m+2*sqrt_p_cos,m-sqrt_p_cos-root_three_sqrt_p_sin,m-sqrt_p_cos+root_three_sqrt_p_sin);
            return DIAGONAL_MATRIX<T,3>(lambda.Sorted().Reversed());
        }
    }
    else{PHYSBAM_FATAL_ERROR();}
}

template<class T> inline void
Jacobi_Transform(const int sweep,const T threshold,T& app,T& apq,T& aqq,T& arp,T& arq,T& v1p,T& v1q,T& v2p,T& v2q,T& v3p,T& v3q)
{
    T epsilon=std::numeric_limits<T>::epsilon();
    T g=100*abs(apq);
    if(sweep > 4 && epsilon*abs(app)>=g && epsilon*abs(aqq)>=g) apq=0;
    else if(abs(apq) > threshold){
        T h=aqq-app,t;
        if(epsilon*abs(h)>=g) t=apq/h;
        else{T theta=(T).5*h/apq;t=1/(abs(theta)+sqrt(1+sqr(theta)));if(theta < 0) t=-t;}
        T cosine=1/sqrt(1+t*t),sine=t*cosine,tau=sine/(1+cosine),da=t*apq;
        app-=da;aqq+=da;apq=0;
        T arp_tmp=arp-sine*(arq+tau*arp);arq=arq+sine*(arp-tau*arq);arp=arp_tmp;
        T v1p_tmp=v1p-sine*(v1q+tau*v1p);v1q=v1q+sine*(v1p-tau*v1q);v1p=v1p_tmp;
        T v2p_tmp=v2p-sine*(v2q+tau*v2p);v2q=v2q+sine*(v2p-tau*v2q);v2p=v2p_tmp;
        T v3p_tmp=v3p-sine*(v3q+tau*v3p);v3q=v3q+sine*(v3p-tau*v3q);v3p=v3p_tmp;}
}

template<class T,int d>
std::pair<DIAGONAL_MATRIX<T,d>,MATRIX<T,d>>
Eigenvalue_Decomposition(const SYMMETRIC_MATRIX<T,d>& M)
{
    typedef VECTOR<T,d> TV;
    if constexpr(d==0) return {};
    else if constexpr(d==1) return {{M.x[0]},{1}};
    else if constexpr(d==2)
    {
        T a=(T).5*M.Trace(),b=(T).5*(M.x[0]-M.x[2]),c=M.x[1];
        T c_squared=sqr(c),m=sqrt(sqr(b)+c_squared),k=M.Determinant();
        TV ev;
        if(a>=0){ev.x=a+m;ev.y=ev.x?k/ev.x:0;}
        else{ev.y=a-m;ev.x=ev.y?k/ev.y:0;}
        TV u=(b>=0?TV(m+b,c):TV(-c,b-m)).Normalized();
        // if order is wrong, matrix is nearly scalar
        return {DIAGONAL_MATRIX<T,2>{ev},{u,u.Perpendicular()}};
    }
    else if constexpr(d==3)
    {
        T a00=M.x[0],a01=M.x[1],a02=M.x[2],a11=M.x[3],a12=M.x[4],a22=M.x[5];
        T v00=1,v01=0,v02=0,v10=0,v11=1,v12=0,v20=0,v21=0,v22=1;
        int sweep;for(sweep=0;sweep<50;sweep++){
            T sum=abs(a01)+abs(a02)+abs(a12);
            if(sum==0) break;
            T threshold=sweep<4?(T)(1./45)*sum:0;
            Jacobi_Transform(sweep,threshold,a00,a01,a11,a02,a12,v00,v01,v10,v11,v20,v21);
            Jacobi_Transform(sweep,threshold,a00,a02,a22,a01,a12,v00,v02,v10,v12,v20,v22);
            Jacobi_Transform(sweep,threshold,a11,a12,a22,a01,a02,v01,v02,v11,v12,v21,v22);}
        assert(sweep<=50);
        return {{a00,a11,a22},{v00,v10,v20,v01,v11,v21,v02,v12,v22}};
    }
    else PHYSBAM_FATAL_ERROR();
}
template std::pair<DIAGONAL_MATRIX<double,1>,MATRIX<double,1,1> > Eigenvalue_Decomposition<double,1>(SYMMETRIC_MATRIX<double,1> const&);
template std::pair<DIAGONAL_MATRIX<double,2>,MATRIX<double,2,2> > Eigenvalue_Decomposition<double,2>(SYMMETRIC_MATRIX<double,2> const&);
template std::pair<DIAGONAL_MATRIX<double,3>,MATRIX<double,3,3> > Eigenvalue_Decomposition<double,3>(SYMMETRIC_MATRIX<double,3> const&);
template std::pair<DIAGONAL_MATRIX<float,1>,MATRIX<float,1,1> > Eigenvalue_Decomposition<float,1>(SYMMETRIC_MATRIX<float,1> const&);
template std::pair<DIAGONAL_MATRIX<float,2>,MATRIX<float,2,2> > Eigenvalue_Decomposition<float,2>(SYMMETRIC_MATRIX<float,2> const&);
template std::pair<DIAGONAL_MATRIX<float,3>,MATRIX<float,3,3> > Eigenvalue_Decomposition<float,3>(SYMMETRIC_MATRIX<float,3> const&);
template DIAGONAL_MATRIX<double,1> Eigenvalues<double,1>(SYMMETRIC_MATRIX<double,1> const&);
template DIAGONAL_MATRIX<double,2> Eigenvalues<double,2>(SYMMETRIC_MATRIX<double,2> const&);
template DIAGONAL_MATRIX<double,3> Eigenvalues<double,3>(SYMMETRIC_MATRIX<double,3> const&);
template DIAGONAL_MATRIX<float,1> Eigenvalues<float,1>(SYMMETRIC_MATRIX<float,1> const&);
template DIAGONAL_MATRIX<float,2> Eigenvalues<float,2>(SYMMETRIC_MATRIX<float,2> const&);
template DIAGONAL_MATRIX<float,3> Eigenvalues<float,3>(SYMMETRIC_MATRIX<float,3> const&);
template DIAGONAL_MATRIX<double,0> Eigenvalues<double,0>(SYMMETRIC_MATRIX<double,0> const&);
template DIAGONAL_MATRIX<double,9> Eigenvalues<double,9>(SYMMETRIC_MATRIX<double,9> const&);
template DIAGONAL_MATRIX<float,0> Eigenvalues<float,0>(SYMMETRIC_MATRIX<float,0> const&);
template DIAGONAL_MATRIX<float,9> Eigenvalues<float,9>(SYMMETRIC_MATRIX<float,9> const&);
template std::pair<DIAGONAL_MATRIX<double,0>,MATRIX<double,0,0> > Eigenvalue_Decomposition<double,0>(SYMMETRIC_MATRIX<double,0> const&);
template std::pair<DIAGONAL_MATRIX<double,9>,MATRIX<double,9,9> > Eigenvalue_Decomposition<double,9>(SYMMETRIC_MATRIX<double,9> const&);
template std::pair<DIAGONAL_MATRIX<float,0>,MATRIX<float,0,0> > Eigenvalue_Decomposition<float,0>(SYMMETRIC_MATRIX<float,0> const&);
template std::pair<DIAGONAL_MATRIX<float,9>,MATRIX<float,9,9> > Eigenvalue_Decomposition<float,9>(SYMMETRIC_MATRIX<float,9> const&);
}
