/* quantum computing emulation library, R. Perry, January 2022
  
supported types:

  state	(Complex)	other real variables (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 State state<Complex,Real>

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

template<typename Real> using State = state<std::complex<Real>,Real>; // with using
State<double> q; // same as state<std::complex<double>,double> q;
*/
#ifndef _QCE_H_
#define _QCE_H_

#include <cmath>
#include <complex>
#include <cstring>	// memcpy(), used in svd()
#include <climits>	// CHAR_BIT
#include <type_traits>  // is_class

#include "lib.h"	// non-template functions, defined in lib.cc

namespace qce
{
//
//   SRAND	set srand() using time(0) or SRAND environment variable
//   drand	0.0 <= drand() < 1.0
//   irand(n)	0 <= irand() < n
//   printb	print binary
//   bitrev	bit reverse, used by qfft()
//   error	print error message and exit
//   emalloc	malloc with exit on error
//
unsigned int SRAND();
double drand();
unsigned long irand( unsigned long n);
void printb( unsigned long x, unsigned int n);
unsigned long bitrev( unsigned long x, unsigned int nbits);
void error( const char *msg);
void *emalloc( unsigned long size);

// template functions:
//
//   abs2	magnitude squared, equivalent to abs(x)**2, or norm(x) for C++ complex
//		(seems to be more accurate than norm())
//   Conj	complex conjugate which returns Real for real argument
//		(std::conj() returns complex<Real>)
//   fft	FFT
//
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); }
template<typename Complex, typename Real> void fft( Complex *X, unsigned long N, int inv);

// print option:
//
//   entanglement measure, decimal state labels, binary state indices and probabilities
//
enum print_option { PRINT_NONE, PRINT_TANGLE, PRINT_LABELS, PRINT_BITS };

// state class and methods:
//
//   state	constructor
//   zero	set state vector to all zeros
//   init	set state k = 1 or random phase
//   swap	swap states i and j
//   tangle	trace distance as a measure of entanglement
//   normalize	scale state vector to unit norm
//   H		Hadamard transformation on one qubit
//   X		NOT transformation on one qubit
//   Z		phase-flip on one qubit
//   CX		controlled NOT on one qubit
//   Hrp	Hadamard with phase shift error
//   qfft	quantum FFT
//   print	display message and qubit state or entanglement measure
//
template <typename Complex = std::complex<double>, typename Real = double> struct state {

  unsigned int n;  // number of bits
  unsigned long N; // number of states, N = 2**n
  Complex *a;      // N amplitudes

  Real eps;	   // machine precision, used in tangle()

 // POSIX constants M_SQRT1_2 and M_PI are only double precision
 // need 1/sqrt(2) and pi here in float, double, or long double precision
 // should change later to use c++20 <numbers> 
  static constexpr Real s2 = std::sqrt((Real)2)/(Real)2;
  static constexpr Real pi = (Real)4*std::atan((Real)1);
  static constexpr Real onethird = (Real)1.0/(Real)3.0; // testing precision

  // allocate a new state vector, uninitialized
  state( unsigned int n) {
    if( n >= CHAR_BIT*sizeof(unsigned long)) error( "state: N overflow");
    this->n = n; N = 1LU << n; unsigned long size = N*sizeof(Complex);
    if( size/sizeof(Complex) != N) error( "state: size overflow");
    a = (Complex*) emalloc(size);
    eps = 1; Real t = 1 + eps/2; while( t > 1) { eps /= 2; t = 1 + eps/2; }
  }

  // destructor
  ~state() { free(a); }

  // set state vector to all zeros
  void zero() {
    for( unsigned long i = 0; i < N; ++i) a[i] = 0;
  }

  // set state k = 1
  void init( unsigned long k) {
    zero(); a[k] = 1;
  }

  // set state k = random phase, Complex must be a complex type
  void init( unsigned long k, int rphase) {
    Complex v = 1;
    if( rphase) { Real a = pi*(2*drand()-1); v = exp(Complex(0,a)); }
    zero(); a[k] = v;
  }

  // swap states i and j
  void swap( unsigned long i, unsigned long j) {
    Complex t = a[i]; a[i] = a[j]; a[j] = t;
  }

  // trace distance as a measure of entanglement
  // distance is 0 for no entanglement, 1 for max entanglement
  //
  // trace distance = 1-fidelity; fidelity = (|<A,B>|/(||A||*||B||))**2
  // where A (B) is the reduced state with qubit k equal to 0 (1)
  //
  // If qubit k is not entangled with any of the other qubits,
  // then the reduced state A with qubit k equal to 0 should be the same
  // as the reduced state B with qubit k equal to 1, i.e. the reduced state
  // should be independent of qubit k.
  //
  // <A,B>/(||A||*||B||)) is the cosine of the angle between vectors A and B.
  // A and B represent the same state if the angle is 0 or 180 degrees,
  // which corresponds to the cosine having a magnitude squared value of 1.
  //
  Real tangle( unsigned int k) const {
    Complex prod = 0; Real A = 0, B = 0, r = 0;
    unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1
    while( i < N) { // i will have k'th bit = 0, j will have k'th bit = 1
      j = i | k1; prod += Conj(a[i])*a[j];
      A += abs2(a[i]); B += abs2(a[j]);
      ++i; i += (i & k1); // skip i values with k'th bit = 1
    }
    // threshold for zero norm squared: 2 * sum { |e + i*e|**2 } = 2 * (2*e*e*(N/2))
    Real t = 2*eps*eps*N;
    if( A > t && B > t) {
      r = 1 - abs2(prod)/(A*B);
      if( r < 0) r = 0; // fixup for roundoff error
    }
    return r;
  }

  // scale state vector to unit norm
  void normalize() {
    Real sum = 0; for( unsigned long i = 0; i < N; ++i) sum += abs2(a[i]);
    sum = std::sqrt(sum); if( sum > 0) for( unsigned long i = 0; i < N; ++i) a[i] /= sum;
  }

  // Hadamard transformation on qubit k
  void H( unsigned int k) {
    if( k >= n) error( "H: bad index");
    unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1
    while( i < N) { // i will have k'th bit = 0, j will have k'th bit = 1
      j = i | k1; Complex t0 = a[i], t1 = a[j];
      a[i] = s2*(t0+t1); a[j] = s2*(t0-t1);
      ++i; i += (i & k1); // skip i values with k'th bit = 1
    }
  }

  // NOT transformation on qubit k
  void X( unsigned int k) {
    if( k >= n) error( "X: bad index");
    unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1
    while( i < N) { // i will have k'th bit = 0, j will have k'th bit = 1
      j = i | k1; swap(i,j);
      ++i; i += (i & k1); // skip i values with k'th bit = 1
    }
  }

  // Z transformation on qubit k
  void Z( unsigned int k) {
    if( k >= n) error( "Z: bad index");
    unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1
    while( i < N) { // i will have k'th bit = 0, j will have k'th bit = 1
      j = i | k1; a[j] = -a[j];
      ++i; i += (i & k1); // skip i values with k'th bit = 1
    }
  }

  // NOT transformation on qubit k controlled by qubit c
  void CX( unsigned int c, unsigned int k) {
    if( c >= n || k >= n || c == k) error( "CX: bad index");
    unsigned long c1 = 1LU << c, k1 = 1LU << k, i = c1, j;
    // printf( "CX: %u %u\n", c, k); // debug
    while( i < N) { // i will have c'th bit = 1, k'th bit = 0
      j = i | k1; // k'th bit = 1
      swap(i,j); // printf( " %lu %lu\n", i, j); // debug
      ++i; i += (i & k1); i |= c1; // skip i values with k'th bit = 1
    }
  }

  // Hadamard transformation on qubit k with phase shift error
  // see http://fog.misty.com/perry/qc/pulse/notes.html
  // Note that Hrp*Hrp = I, and rp = p/(pi/2) would be 1 for no error
  // Complex must be a complex type
  void Hrp( unsigned int k, Real rp) {
    Real C, S; if( k >= n) error( "H: bad index");
    if( rp == 1.0) { C = 0; S = 1; } else { Real p = rp*(pi/2); C = cos(p), S = sin(p); }
    Complex c01(S,-C), c10(S,C);
    unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1
    while( i < N) { // i will have k'th bit = 0, j will have k'th bit = 1
      j = i | k1; Complex t0 = a[i], t1 = a[j];
      a[i] = s2*(t0 + c01*t1); a[j] = s2*(c10*t0 - t1);
      ++i; i += (i & k1); // skip i values with k'th bit = 1
    }
  }

  // quantum FFT, Complex must be a complex type
  void qfft( int inv = 0) {
    fft<Complex,Real>( a, N, inv); Real s = 1/std::sqrt((Real)N);
   // scale values and bit-reverse the indices
   // see https://stackoverflow.com/questions/932079/in-place-bit-reversed-shuffle-on-an-array
    for( unsigned long i = 0, j = 0; i < N; ++i) {
      j = bitrev( i, n);
      if( i < j) { swap( i, j); a[i] *= s; a[j] *= s; }
      else if( i == j) a[i] *= s;
    }
  }

  // display message and qubit state or entanglement measure
  void print( const char *msg, print_option option = PRINT_NONE) const {
    static const Complex z = 0.0;
    std::cout << msg << ":" << ((option == PRINT_BITS) ? "\n" : "");
    if( option == PRINT_TANGLE) {
      for( unsigned int k = 0; k < n; ++k) std::cout << " " << tangle(k);
    } else for( unsigned long i = 0; i < N; ++i) {
      if( option == PRINT_BITS) {
        std::cout << "  "; printb( i, n); std::cout << " "; int prec = std::cout.precision(4);
	std::cout.width(8); std::cout << std::left << abs2(a[i]) << std::right;
        std::cout.precision(prec); std::cout << " "; // use parens for Real state:
	if( !std::is_class<Complex>::value) std::cout << "(";
	std::cout << z+a[i]; // z+ to make -0 print as just 0
	if( !std::is_class<Complex>::value) std::cout << ")";
	std::cout << "\n";
      } else {
	std::cout << " ";
        if( option == PRINT_LABELS) std::cout << i << ":";
	if( !std::is_class<Complex>::value) std::cout << "(";
	std::cout << a[i]; if( !std::is_class<Complex>::value) std::cout << ")";
      }
    }
    if( option != PRINT_BITS) std::cout << '\n';
  }

};

// << overload
template<typename Complex, typename Real>
std::ostream& operator<<(std::ostream& s, const state<Complex,Real>& q) {
  for( unsigned long i = 0; i < q.N; ++i) { // use parens for Real state
    s << " ";    if( !std::is_class<Complex>::value) std::cout << "(";
    s << q.a[i]; if( !std::is_class<Complex>::value) std::cout << ")";
  }
  return s << "\n";
}

#include "fft.h"
#include "svd.h"
#include "simplex.h"

} // namespace qce

#endif // _QCE_H_