//#####################################################################
// Copyright 2022-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/Math_Tools/minmax.h>
#include <Core/Math_Tools/sqr.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/MATRIX_MXN.h>
#include <Core/Matrices/ROTATION.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
namespace PhysBAM{
namespace SVD3{
template<class T>
VECTOR<T,2> Givens(T a,T b)
{
    return VECTOR<T,2>(a,-b).Normalized();
}
// Apply Givens rotation to matrix (on the left)
template<class T>
void Apply_L(MATRIX<T,3>& A,VECTOR<T,2> u,int i,int j)
{
    auto v=A.Row(i),w=A.Row(j);
    A.Set_Row(i,u.x*v+u.y*w);
    A.Set_Row(j,u.x*w-u.y*v);
}
// Apply Givens rotation to matrix (on the right)
template<class T>
void Apply_R(MATRIX<T,3>& A,VECTOR<T,2> u,int i,int j)
{
    auto v=A.Column(i),w=A.Column(j);
    A.Set_Column(i,u.x*v+u.y*w);
    A.Set_Column(j,u.x*w-u.y*v);
}
// Puts a zero in A(r1,c)
template<class T>
void Zero_L(MATRIX<T,3>& U,MATRIX<T,3>& A,int r0,int r1,int c)
{
    auto g=Givens(A(r0,c),A(r1,c));
    Apply_L(A,g,r1,r0);
    Apply_R(U,g,r1,r0);
    A(r1,c)=0;
}
// Puts a zero in A(r,c1)
template<class T>
void Zero_R(MATRIX<T,3>& V,MATRIX<T,3>& A,int c0,int c1,int r)
{
    auto g=Givens(A(r,c0),A(r,c1));
    Apply_R(A,g,c1,c0);
    Apply_R(V,g,c1,c0);
    A(r,c1)=0;
}
template<class T>
void Zero_Chasing(MATRIX<T,3>& U,MATRIX<T,3>& A,MATRIX<T,3>& V)
{
    // Initial sparsity:  [ * * * ; * * * ; 0 * * ]
    Zero_L(U,A,0,1,0); // [ * * * ; 0 * * ; 0 * * ]
    Zero_R(V,A,1,2,0); // [ * * 0 ; 0 * * ; 0 * * ]
    Zero_L(U,A,1,2,1); // [ * * 0 ; 0 * * ; 0 0 * ]
}
template<class T>
void Swap_Col(MATRIX<T,3>& U,int i,int j)
{
    auto u=U.Column(i),v=U.Column(j);
    U.Set_Column(i,v);
    U.Set_Column(j,-u);
}
template<class T>
void Swap_Row(MATRIX<T,3>& U,int i,int j)
{
    auto u=U.Row(i),v=U.Row(j);
    U.Set_Row(i,v);
    U.Set_Row(j,-u);
}
template<class T>
void Swap(MATRIX<T,3>& U,DIAGONAL_MATRIX<T,3>& D,MATRIX<T,3>& V,int i,int j)
{
    Swap_Col(U,i,j);
    Swap_Col(V,i,j);
    std::swap(D.x(i),D.x(j));
}
template<class T>
void Flip_Sign(MATRIX<T,3>& U,DIAGONAL_MATRIX<T,3>& D,int i,int j)
{
    D.x(i)=-D.x(i);
    D.x(j)=-D.x(j);
    U.Set_Column(i,-U.Column(i));
    U.Set_Column(j,-U.Column(j));
}
template<class T>
void Sort_SV(MATRIX<T,3>& U,DIAGONAL_MATRIX<T,3>& D,MATRIX<T,3>& V)
{
    int i=abs(D.x).Arg_Max();
    if(i>0) Swap(U,D,V,0,i);
    if(abs(D.x.y)<abs(D.x.z)) Swap(U,D,V,1,2);
    if(D.x.x<0) Flip_Sign(U,D,0,D.x.y<0?1:2);
    else if(D.x.y<0) Flip_Sign(U,D,1,2);
}
template<class T>
void Top_Block(MATRIX<T,3>& U,MATRIX<T,3>& B,MATRIX<T,3>& V,DIAGONAL_MATRIX<T,3>& D)
{
    MATRIX<T,2> a(B(0,0),B(1,0),B(0,1),B(1,1));
    auto [u,d,v]=a.Singular_Value_Decomposition();
    Apply_R(U,u.Column(0),0,1);
    Apply_R(V,v.Column(0),0,1);
    D.x.x=d.x.x;
    D.x.y=d.x.y;
    D.x.z=B(2,2);
    Sort_SV(U,D,V);
}
template<class T>
void Bot_Block(MATRIX<T,3>& U,MATRIX<T,3>& B,MATRIX<T,3>& V,DIAGONAL_MATRIX<T,3>& D)
{
    MATRIX<T,2> a(B(1,1),B(2,1),B(1,2),B(2,2));
    auto [u,d,v]=a.Singular_Value_Decomposition();
    Apply_R(U,u.Column(0),1,2);
    Apply_R(V,v.Column(0),1,2);
    D.x.x=B(0,0);
    D.x.y=d.x.x;
    D.x.z=d.x.y;
    Sort_SV(U,D,V);
}
template<class T>
bool Check_Tolerances(MATRIX<T,3>& B,MATRIX<T,3>& U,DIAGONAL_MATRIX<T,3>& D,MATRIX<T,3>& V,T tau)
{
    if(abs(B(1,2))<=tau){
        Top_Block(U,B,V,D);
        return true;}
    if(abs(B(0,1))<=tau){
        Bot_Block(U,B,V,D);
        return true;}
    if(abs(B(1,1))<=tau){
        auto g=Givens(B(1,2),B(2,2)).Orthogonal_Vector();
        Apply_L(B,g,2,1);
        Apply_R(U,g,2,1);
        Top_Block(U,B,V,D);
        return true;}
    if(abs(B(2,2))<=tau){
        Zero_R(V,B,1,2,1);
        Zero_R(V,B,0,2,0);
        Top_Block(U,B,V,D);
        return true;}
    if(abs(B(0,0))<=tau){
        Zero_L(U,B,1,0,1);
        Zero_L(U,B,2,0,2);
        Bot_Block(U,B,V,D);
        return true;}
    return false;
}

template<class T> bool
Lazy_Sort(MATRIX<T,3>& U,MATRIX<T,3>& B,MATRIX<T,3>& V,int i,int j)
{
    T flip_ratio=100;
    if(abs(B(i,i))*flip_ratio>=abs(B(j,j))) return false;
    Swap_Row(B,i,j);
    Swap_Col(U,i,j);
    Swap_Col(B,i,j);
    Swap_Col(V,i,j);
    return true;
}

template<class T> void
Bidiagonalize(MATRIX<T,3>& U,MATRIX<T,3>& B,MATRIX<T,3>& V)
{
    Zero_L(U,B,1,2,0); // [ * * * ; * * * ; 0 * * ]
    Zero_Chasing(U,B,V); // [ * * 0 ; 0 * * ; 0 0 * ]
}
//#####################################################################
// Function Singular_Value_Decomposition
//#####################################################################
// U and V rotations, smallest singular value possibly negative
// @techreport{gast:2016:qrsvd,
//   title={Implicit-shifted Symmetric QR Singular Value Decomposition of 3x3 Matrices},
//   author={Gast, T. and Fu, C. and Jiang, C. and Teran, J.},
//   year={2016},
//   institution={University of California Los Angeles}
// }
template<class T> std::tuple<MATRIX<T,3>,DIAGONAL_MATRIX<T,3>,MATRIX<T,3>>
Singular_Value_Decomposition(MATRIX<T,3> B)
{
    DIAGONAL_MATRIX<T,3> D;
    MATRIX<T,3> U,V;
    U+=1;
    V+=1;
    Bidiagonalize(U,B,V);

    T nu=16*std::numeric_limits<T>::epsilon();
    T tau=nu*max((T).5*B.Frobenius_Norm(),(T)1);
    for(int it=0;;it++)
    {
        if(Check_Tolerances(B,U,D,V,tau)) return {U,D,V};
        if(Lazy_Sort(U,B,V,0,1) | Lazy_Sort(U,B,V,0,2) | Lazy_Sort(U,B,V,1,2)){
            Bidiagonalize(U,B,V);
            if(Check_Tolerances(B,U,D,V,tau)) return {U,D,V};}

        T al0=B(0,0),al1=B(1,1),al2=B(2,2);
        T be0=B(0,1),be1=B(1,2);
        T ga0=al0*be0,ga1=al1*be1;
        T a0=al1*al1+be0*be0;
        T a1=al2*al2+be1*be1;
        T b0=ga1;
        T d=(T).5*(a0-a1);
        T mu=a1-b0*b0/(d+sign_nonzero(d)*sqrt(d*d+b0*b0));

        auto g=Givens(al0*al0-mu,ga0);
        Apply_R(B,g,1,0);
        Apply_R(V,g,1,0);
        Zero_Chasing(U,B,V);
    }
}
}
namespace SVD2{
template<class T> std::tuple<MATRIX<T,2>,SYMMETRIC_MATRIX<T,2>>
Indefinite_Polar_Decomposition(const MATRIX<T,2>& M)
{
    T tr=M.Trace(),cosine,sine;
    if(tr==0){cosine=0;sine=1;}
    else
    {
        T t=(M.x[1]-M.x[2])/tr;
        cosine=1/sqrt(1+t*t);
        sine=t*cosine;
    }
    MATRIX<T,2> Q{cosine,sine,-sine,cosine};
    return {Q,Transpose_Times_With_Symmetric_Result(Q,M)};
}
template<class T> std::tuple<MATRIX<T,2>,DIAGONAL_MATRIX<T,2>,MATRIX<T,2>>
Singular_Value_Decomposition(const MATRIX<T,2>& M)
{
    auto [Q,S]=Indefinite_Polar_Decomposition(M);
    auto [ev,V]=S.Solve_Eigenproblem();
    if(ev.x.y<0 && abs(ev.x.y)>=abs(ev.x.x))
    {
        ev=DIAGONAL_MATRIX<T,2>(-ev.x.y,-ev.x.x);
        Q=-Q;
        V=MATRIX<T,2>(V.x[2],V.x[3],-V.x[0],-V.x[1]);
    }
    return {Q*V,ev,V};
}
}
namespace SVD1{
template<class T> std::tuple<MATRIX<T,1>,DIAGONAL_MATRIX<T,1>,MATRIX<T,1>>
Singular_Value_Decomposition(const MATRIX<T,1>& M)
{
    return {{1},{M.x[0]},{1}};
}
}
template<class T,int d>
std::tuple<MATRIX<T,d>,DIAGONAL_MATRIX<T,d>,MATRIX<T,d>>
Singular_Value_Decomposition(const MATRIX<T,d>& A)
{
    if constexpr(d==0) return {};
    else if constexpr(d==1) return SVD1::Singular_Value_Decomposition(A);
    else if constexpr(d==2) return SVD2::Singular_Value_Decomposition(A);
    else return SVD3::Singular_Value_Decomposition(A);
}
template<class T>
std::tuple<MATRIX<T,3,2>,DIAGONAL_MATRIX<T,2>,MATRIX<T,2>>
Singular_Value_Decomposition(const MATRIX<T,3,2>& A)
{
    if constexpr(is_same<T,float>::value)
    {
        auto [U,S,V]=MATRIX<double,3,2>(A).Singular_Value_Decomposition();
        return {MATRIX<T,3,2>(U),DIAGONAL_MATRIX<T,2>(S),MATRIX<T,2>(V)};
    }
    else
    {
        auto [lambda,V]=A.Normal_Equations_Matrix().Solve_Eigenproblem();
        if(lambda.x.y<0) lambda.x=Componentwise_Max(lambda.x,(T)0);
        auto u=(A*V.Column(0)).Normalized();
        auto v=Weighted_Normal(A).Cross(u);
        if(!v.Normalize())
            v=u.Unit_Orthogonal_Vector();
        return {{u,v},DIAGONAL_MATRIX<T,2>{sqrt(lambda.x)},V};
    }
}
template std::tuple<MATRIX<double,1,1>,DIAGONAL_MATRIX<double,1>,MATRIX<double,1,1> > Singular_Value_Decomposition<double,1>(MATRIX<double,1,1> const&);
template std::tuple<MATRIX<double,2,2>,DIAGONAL_MATRIX<double,2>,MATRIX<double,2,2> > Singular_Value_Decomposition<double,2>(MATRIX<double,2,2> const&);
template std::tuple<MATRIX<double,3,2>,DIAGONAL_MATRIX<double,2>,MATRIX<double,2,2> > Singular_Value_Decomposition<double>(MATRIX<double,3,2> const&);
template std::tuple<MATRIX<double,3,3>,DIAGONAL_MATRIX<double,3>,MATRIX<double,3,3> > Singular_Value_Decomposition<double,3>(MATRIX<double,3,3> const&);
template std::tuple<MATRIX<float,1,1>,DIAGONAL_MATRIX<float,1>,MATRIX<float,1,1> > Singular_Value_Decomposition<float,1>(MATRIX<float,1,1> const&);
template std::tuple<MATRIX<float,2,2>,DIAGONAL_MATRIX<float,2>,MATRIX<float,2,2> > Singular_Value_Decomposition<float,2>(MATRIX<float,2,2> const&);
template std::tuple<MATRIX<float,3,2>,DIAGONAL_MATRIX<float,2>,MATRIX<float,2,2> > Singular_Value_Decomposition<float>(MATRIX<float,3,2> const&);
template std::tuple<MATRIX<float,3,3>,DIAGONAL_MATRIX<float,3>,MATRIX<float,3,3> > Singular_Value_Decomposition<float,3>(MATRIX<float,3,3> const&);
template std::tuple<MATRIX<double,0,0>,DIAGONAL_MATRIX<double,0>,MATRIX<double,0,0> > Singular_Value_Decomposition<double,0>(MATRIX<double,0,0> const&);
template std::tuple<MATRIX<float,0,0>,DIAGONAL_MATRIX<float,0>,MATRIX<float,0,0> > Singular_Value_Decomposition<float,0>(MATRIX<float,0,0> const&);
}
