/* Singular Value Decomposition and associated matrix operations

R. Perry, 10 Jan 2022, https://fog.misty.com/perry/svd/svd.html

Nash SVD algorithm adapted for complex data and row operations instead of columns.
 Reference: Compact Numerical Methods for Computers, second edition, J. C. Nash,
 Adam Hilger publishing, 1990, ISBN 0-85274-318-1

Supported types:

  Complex		Real
  ---------------	---------------------------
  float			float
  double		double
  long double		long double
  complex<float>	float
  complex<double>	double
  complex<long double>	long double

// recommended usage:
#define Real double		    // or float or long double
#define Complex std::complex<Real>  // or just Real
#define SVD svd<Complex,Real>

// other examples:
#define SVD svd<std::complex<double>,double>
#define SVD svd			    // same, using defaults
#define SVD svd<float,float>	    // real, float precision

template<typename Real> using SVD = svd<std::complex<Real>,Real>; // with using
SVD<double> q; // same as svd<std::complex<double>,double> q;
*/
#ifndef _SVD_H_
#define _SVD_H_

#include <cmath>
#include <complex>

#ifdef SVD_STANDALONE // if not being compiled as part of qce++
#include <cstring> // memcpy()
namespace Svd
{
#endif

// non-template functions, defined in svd.cc:
//
//   SRAND	set srand() using time(0) or SRAND environment variable
//   drand	0.0 <= drand() < 1.0
//   error	print error message and exit
//   emalloc	malloc with exit on error
//
#ifdef SVD_STANDALONE // these are provided in qce++ lib.cc
unsigned int SRAND();
double drand();
void error( const char *msg);
void *emalloc( unsigned long size);
#endif

// template functions:
//
//   abs2	magnitude squared, equivalent to abs(x)**2, or norm(x) for C++ complex,
//		 |x|**2 for vectors
//   Conj	complex conjugate which returns Real for real argument
//		 (std::conj() returns complex<Real>)
//   matdiff	difference check, |A-AA|
//   matprint	print matrix in octave/matlab format
//   prod	vector inner product
//   rotate	plane rotation
//   ucheck	unitary check, |Q*Q'-I| should be close to 0
//
#ifdef SVD_STANDALONE // these are provided in qce++ qce.h
template<typename Real> Real abs2(Real x) {
  return x*x;
}
template<typename Real> Real abs2(std::complex<Real> x) {
  Real r=real(x), i=imag(x); return r*r+i*i;
}
template<typename Real> Real Conj(Real x) {
  return x;
}
template<typename Real> std::complex<Real> Conj(std::complex<Real> x) {
  return std::conj(x);
}
#endif

template<typename Real> Real abs2(const Real *x, unsigned long m) {
  Real r = 0; while( m--) { r += abs2( *x++); } return r;
}
template<typename Real> Real abs2(const std::complex<Real> *x, unsigned long m) {
  Real r = 0; while( m--) { r += abs2( *x++); } return r;
}

template<typename Complex, typename Real> Real matdiff(
 const Complex *A, const Complex *AA, unsigned long n, unsigned long m, Real eps) {
 // eps not used, just there so template knows what Real is
  unsigned long j, k; Real r = 0; if( A && AA)
   for( j = 0; j < n; ++j) for( k = 0; k < m; ++k) r += abs2( *A++ - *AA++);
  return sqrt(r);
}
template<typename Real> void matprint( std::ostream& s,
 const char *msg, const Real *V, unsigned long n, unsigned long m) {
  int prec = s.precision( std::numeric_limits<Real>::digits10); // set full precision
  unsigned long j, k; s << msg << " = [\n"; if( V)
   for( j = 0; j < n; ++j) { for( k = 0; k < m; ++k, ++V) { s << " " << *V; } s << "\n"; }
  s << "];\n"; s.precision( prec); // restore default precision
}
template<typename Real> void matprint( std::ostream& s,
 const char *msg, const std::complex<Real> *V, unsigned long n, unsigned long m) {
  int prec = s.precision( std::numeric_limits<Real>::digits10); // set full precision
  unsigned long j, k; Real r, i; s << msg << " = [\n"; if( V)
   for( j = 0; j < n; ++j) { for( k = 0; k < m; ++k, ++V) { r = real(*V); i = imag(*V);
    s << " " << r << ((i < 0) ? "-" : "+") << fabs(i) << "i"; } s << "\n"; }
  s << "];\n"; s.precision( prec); // restore default precision
}
template<typename Complex> Complex prod( const Complex *x, const Complex *y, unsigned long m) {
  Complex r = 0; while( m--) r += Conj(*x++) * *y++;
  return r;
}
template<typename Complex, typename Real> void rotate(
 Complex *Vj, Complex *Vk, unsigned long m, Real c, Complex s, int swap) {
  for( unsigned long i = 0; i < m; ++i) { Complex vj = *Vj, vk = *Vk;
    if( swap) { *Vj++ = s*vj + c*vk; *Vk++ = c*vj - Conj(s)*vk; }
    else      { *Vj++ = c*vj + s*vk; *Vk++ = c*vk - Conj(s)*vj; } }
}
template<typename Complex, typename Real> Real ucheck(
 const Complex *V, unsigned long n, unsigned long m, Real eps) {
 // eps not used, just there so template knows what Real is
  unsigned long j, k; const Complex *Vj, *Vk; Real d, r = 0;
  for( j = 0, Vj = V; j < n; ++j, Vj += m) { d = abs2(Vj,m)-1; r += d*d; }
  for( j = 0, Vj = V; j < n-1; ++j, Vj += m)
    for( k = j+1, Vk = Vj+m; k < n; ++k, Vk += m) r += 2*abs2(prod(Vj,Vk,m));
  return sqrt(r);
}

// svd class and methods:
//
//   svd	constructor, performs the decomposition
//   ~svd	destructor
//   check	compute AA = Q'*S*V = Q(1:rank,:)'*diag(S,rank,rank)*V(1:rank,:)
//
template <typename Complex = std::complex<double>, typename Real = double> struct svd {

 // Multiply by a sequence of plane rotations on the left to orthogonalize the rows of A:
 // Q*[A|I] = [B|Q], then factor out the row norms of B: Q*A=S*V, yielding A=Q'*S*V
 // with Q and V unitary, S diagonal and real. Sizes: A=(n,m), Q=(n,n), S=(n), V=(n,m)
 // All matrices are stored by rows in 1D arrays.
 //
  int copy; unsigned long n, m, rank, count = 0 /* sweep count */;
  Complex *AA = 0, *Q, *V; Real *S, eps /* machine precision */;

  // perform SVD: V overwrites A if copy is zero; no sweep count limit if limit is zero
  svd( Complex *A, unsigned long n, unsigned long m, int copy = 1, unsigned long limit = 20) {
    this->n = n; this->m = m; this->copy = copy; V = A;
    if( copy) { unsigned long size = n*m*sizeof(Complex);
      V = (Complex*) emalloc(size); std::memcpy( V, A, size); }
    Q = (Complex*) emalloc(n*n*sizeof(Complex)); S = (Real*) emalloc(n*sizeof(Real));
    unsigned long j, k; Complex *Qj, *Qk, *Vj, *Vk;
    for( j = 0, Qj = Q; j < n; ++j, Qj += n) { // Q = I
      for( k = 0; k < n; ++k) Qj[k] = 0;
      Qj[j] = 1; }
    for( j = 0, Vj = V; j < n; ++j, Vj += m) S[j] = abs2( Vj, m);
    eps = 1; Real t = 1 + eps/2; while( t > 1) { eps /= 2; t = 1 + eps/2; }
    Real tol = m*eps*eps /* tolerance for zero */, q, r, qr, v, c; Complex p, s;
    int done, swap; do { done = 1; ++count; // do until no rotations or sweep limit reached
      for( j = 0, Vj = V, Qj = Q; j < n-1; ++j, Vj += m, Qj += n)
        for( k = j+1, Vk = Vj+m, Qk = Qj+n; k < n; ++k, Vk += m, Qk += n) {
	  // make rows j & k of V orthogonal
          q = S[j]; r = S[k]; swap = 0;
	  qr = q*r; if( qr < tol) continue; // if norms are too small then skip
	  p = prod(Vj,Vk,m); if( abs2(p) < tol*qr) continue; // if orthogonal then skip
	  if( q < r) { // rows out of order, incorporate swap with rotate
	    swap = 1; Real t = q; q = r; r = t; p = Conj(p); }
	  // std::cout << q << " " << r << " " << p << " "; // debug
	  // std::cout << "jk: " << j << " " << k << " " << swap << '\n'; // debug
	  done = 0; p = p/q; r = 1-r/q; v = sqrt(4*abs2(p) + r*r);
	  c = sqrt((1+r/v)/2); s = Conj(p)/(v*c);
	  // std::cout << c << " " << s << "\n"; // debug
	  rotate( Vj, Vk, m, c, s, swap); rotate( Qj, Qk, n, c, s, swap);
	  S[j] = abs2( Vj, m); S[k] = abs2( Vk, m); } // update norms
    } while( !done && (!limit || count < limit));
    if( !done) std::cerr << "svd: reached sweep limit (" << limit << ")" << std::endl;
    // find rank and normalize
    rank = 0; if( S[0] == 0) return;
    tol *= m*S[0]; // tol=(m*eps*eps)*m*S[0] // relative tolerance for squares of singular values
    for( j = 0, Vj = V; j < n; ++j, Vj += m) { if( S[j] > tol) ++rank; else break;
     S[j] = sqrt(S[j]); for( k = 0; k < m; ++k) Vj[k] /= S[j]; }
  }

  ~svd() { free(Q); free(S); free(AA); if( copy) free(V); }

  void check() { 
    unsigned long i, j, k; Complex qc, *AAjk, *Qi, *Vi;
    if( !AA) AA = (Complex*) emalloc( n*m*sizeof(Complex));
    for( j = 0, AAjk = AA; j < n; ++j) for( k = 0; k < m; ++k) *AAjk++ = 0;
    for( i = 0, Qi = Q, Vi = V; i < rank; ++i, Qi += n, Vi += m)
      for( j = 0, AAjk = AA; j < n; ++j) { qc = Conj(Qi[j])*S[i];
        for( k = 0; k < m; ++k) *AAjk++ += qc*Vi[k]; }
  }

};

// << overload
template<typename Complex, typename Real>
std::ostream& operator<<(std::ostream& s, const svd<Complex,Real>& a) {
  s << "n = " << a.n << "; m = " << a.m << "; r = " << a.rank << ";\n";
  matprint( s, "Q", a.Q, a.n, a.n); matprint( s, "V", a.V, a.n, a.m);
  matprint( s, "S", a.S, 1, a.n); if( a.AA) matprint( s, "AA", a.AA, a.n, a.m);
  return s;
}

#ifdef SVD_STANDALONE // if not being compiled as part of qce++
} // namespace Svd
#endif

#endif // _SVD_H_