// quantum computing emulation library
//
// R. Perry, Jan. 2018
//
#include <stdio.h>
#include <tgmath.h>
#include <stdlib.h> // rand(), srand()
#include <string.h> // memcpy()
#include <limits.h> // CHAR_BIT
#include <float.h>  // DBL_EPSILON
#include <time.h>   // time()
#include "qce.h"

//------------------------------ private ------------------------------
//
//  t		temporary state copy used in operate()
//  tcopy	make temporary state copy
//  bitrev	bit reverse
//  pi, s2	pi and sqrt(2)/2 constants
//  _init	set pi and s2
//
static State t;
//
static void tcopy( State q)
{
  if( !t.a) t = Q(q.n); else if( t.n != q.n) { free(t.a); t = Q(q.n); }

  memcpy( t.a, q.a, q.N*sizeof(double complex));
}
//
static unsigned long bitrev( unsigned long x, unsigned int nbits)
{
  unsigned long r = 0;

  for( unsigned int i = 0; i < nbits; ++i) { r <<= 1; r |= x & 1; x >>= 1; }

  return r;
}
//
static double pi = 0, s2 = 0;
//
static void _init( void)
{
  if( pi == 0) { pi = 4*atan(1); s2 = sqrt(2)/2; }
}

// if C11 CMPLX macro is not available
// https://gcc.gnu.org/onlinedocs/gcc-8.2.0/gcc/Other-Builtins.html#index-_005f_005fbuiltin_005fcomplex
//
#ifndef CMPLX
#define CMPLX(x,y) __builtin_complex((x),(y))
#endif

//------------------------------ public ------------------------------
//
// set srand() using time(0) or SRAND environment variable
//
unsigned int SRAND( void)
{
  char *env = getenv( "SRAND");  unsigned int seed = 0;  if( env) seed = atoi(env);

  if( seed == 0) seed = time(0);

  srand( seed);  return seed;
}

// allocate a new state vector, uninitialized
//
State Q(int n)
{
  if( n >= CHAR_BIT*sizeof(unsigned long)) error( "Q: N overflow");

  State q;  q.n = n;  q.N = 1LU << n;  unsigned long size = q.N*sizeof(double complex);

  if( size/sizeof(double complex) != q.N) error( "Q: size overflow");

  q.a = emalloc(size);  return q;
}

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

// initialize state vector, all zeros except q.a[k] = 1 * optional_random_phase
//
void init( State q, unsigned long k, int rphase)
{
  _init();  double complex v = 1;

  if( rphase) { double a = pi*(2*DRAND-1); v = exp(CMPLX(0,a)); }

  zero(q);  q.a[k] = v;
}

// swap states i and j
//
void swap( State q, unsigned long i, unsigned long j)
{
  double complex t = q.a[i];  q.a[i] = q.a[j];  q.a[j] = t;
}

// tensor product
//
State product( State r, State s)
{
  State q = Q( r.n + s.n);

  unsigned long i, j, k = 0;

  for( i = 0; i < r.N; ++i) for( j = 0; j < s.N; ++j) q.a[k++] = r.a[i]*s.a[j];

  return q;
}

// 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.
//
double tangle( State q, unsigned int k)
{
  double complex prod = 0;  double A = 0, B = 0, r = 0;

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  prod += conj(q.a[i])*q.a[j];  A += abs2(q.a[i]);  B += abs2(q.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*(q.N/2))
 //
  double t = 2*DBL_EPSILON*DBL_EPSILON*q.N;

  if( A > t && B > t) r = 1 - abs2(prod)/(A*B);

  return r;
}

// Hadamard transformation on qubit k
//
void H( State q, unsigned int k)
{
  _init();  if( k >= q.n) error( "H: bad index");

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  double complex t0 = q.a[i], t1 = q.a[j];

    q.a[i] = s2*(t0+t1);  q.a[j] = s2*(t0-t1);

    ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }
}

#if 0 // algorithm from https://arxiv.org/abs/1801.01037 and https://arxiv.org/abs/1601.07195
void H( State q, unsigned int k)
{
  _init();  if( k >= q.n) error( "H: bad index");
  unsigned long k2 = 1LU << k, k22 = k2 << 1; // 2**k, 2**(k+1)
  for( unsigned long i = 0; i < q.N; i += k22) {
    unsigned long ik2 = i + k2;
    for( unsigned long j = i, jk2 = ik2; j < ik2; ++j, ++jk2) {
      double complex t0 = q.a[j], t1 = q.a[jk2];
      q.a[j] = s2*(t0+t1);  q.a[jk2] = s2*(t0-t1);
    }
  }
}
#endif

// Hadamard transformation on qubit k with non-resonant pulse error
// see http://fog.misty.com/perry/qc/pulse/notes.html
// ra = a/d would be 0 for no error
//
void Hra( State q, unsigned int k, double ra)
{
  _init();  if( k >= q.n) error( "H: bad index");

  double rb = sqrt(1-ra*ra), A = pi/(4*rb), C = cos(A), S = sin(A);

  double raS = ra*S, rbS = rb*S;

  double complex c00 = CMPLX(C,-raS), c11 = CMPLX(C,raS);

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  double complex t0 = q.a[i], t1 = q.a[j];

    q.a[i] = c00*t0 + rbS*t1;  q.a[j] = rbS*t0 - c11*t1;

    ++i;  i += (i & k1); // 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 also that Hrp*Hrp = I
// rp = p/(pi/2) would be 1 for no error
//
void Hrp( State q, unsigned int k, double rp)
{
  _init();  if( k >= q.n) error( "H: bad index");

  double p = rp*(pi/2), C = cos(p), S = sin(p);

 // printf( "C = %g, 1-C = %g, S = %g, 1-S = %g\n", C, 1-C, S, 1-S);
 // with rp=1.0: C = 6.12323e-17, 1-C = 1, S = 1, 1-S = 0
 // use more accurate values for this case:
 //
  if( rp == 1.0) { C = 0; S = 1; }

  double complex c01 = CMPLX(S,-C), c10 = CMPLX(S,C);

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  double complex t0 = q.a[i], t1 = q.a[j];

    q.a[i] = s2*(t0 + c01*t1);  q.a[j] = s2*(c10*t0 - t1);

    ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }
}

// NOT transformation on qubit k
//
void X( State q, unsigned int k)
{
  if( k >= q.n) error( "X: bad index");

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  swap( q, i, j);

    ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }
}

// Z transformation on qubit k
//
void Z( State q, unsigned int k)
{
  if( k >= q.n) error( "Z: bad index");

  unsigned long i = 0, j, k1 = 1LU << k; // k1 = 0...010...0, k'th bit = 1

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  q.a[j] = -q.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( State q, unsigned int c, unsigned int k)
{
  if( c >= q.n || k >= q.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 < q.N)  // i will have c'th bit = 1, k'th bit = 0
  {
    j = i | k1; // k'th bit = 1

    swap( q, i, j);  // printf( " %lu %lu\n", i, j); // debug

    ++i;  i += (i & k1);  i |= c1; // skip i values with k'th bit = 1
  }
}

// measurement
//
unsigned long M( State q)
{
  double P = DRAND, R = 0; // R will be cumulative probability

  for( unsigned long i = 0; i < q.N; ++i)
  {
    R += abs2(q.a[i]);  if( P < R) return i;
  }

 // R should be 1.0 at this point, and P < 1.0, so should not get here.
 // But could happen if R < 1.0 due to accumulation of roundoff errors.
 // error( "M: measurement failed");
 // fprintf( stderr, "qce: WARNING, measurement failed, P = %g, R = %g\n", P, R);

  return q.N-1;
}

// reduce state to the top n qubits based on a measurement m
//
State reduce( State q, unsigned int n, unsigned long m)
{
  if( n >= q.n) error( "reduce: n >= q.n");

  State r = Q(n);  unsigned long k = q.n-n, mask = (1LU<<k)-1;

 // m is measured value from the bottom k qubits
 //
  if( m > mask) error( "reduce: measurement out of range");

  double s = 0;

  for( unsigned long i = 0, j = 0; i < r.N; ++i) // efficient traverse r.N
  {
    j = (i << k) | m;  r.a[i] = q.a[j];  s += abs2(q.a[j]);

#if 0 // debug
      printf( "reduce: i = %lu, j = %lu, q.a[j] = (%g,%g)\n", i, j, creal(q.a[j]), cimag(q.a[j]) );
#endif
  }

  if( s == 0) error( "reduce: magnitude is 0");

 // normalize
 //
  s = sqrt(s);  for( unsigned long i = 0; i < r.N; ++i) r.a[i] /= s;

  return r;
}

// collapse state based on a measurement m of qubit k
//
void collapse( State q, unsigned int k, unsigned int m)
{
  if( k >= q.n) error( "collapse: k >= q.n");

  if( m > 1) error( "collapse: m > 1"); // m must be 0 or 1

  double s = 0;  unsigned long i = 0, j, k1 = 1LU << k;

  while( i < q.N) // i will have k'th bit = 0, j will have k'th bit = 1
  {
    j = i | k1;  s += abs2(q.a[m?j:i]);  q.a[m?i:j] = 0;

    ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }

  if( s == 0) error( "reduce: magnitude is 0");

 // normalize
 //
  s = sqrt(s);  for( i = 0; i < q.N; ++i) q.a[i] /= s;
}

// q = (a,b) -> (a,f(a,b,z))
//
// for an operation which corresponds to a permutation of the state
// where a represents the top n bits and b the bottom m bits with q.n = n + m
//
void operate( State q, unsigned int n, Func f, unsigned long z)
{
  unsigned int m = q.n - n;
  unsigned long M = 1LU << m, mask = M-1, a, b, c, X, Y;

  if( n >= q.n) error( "operate: n >= q.n");

  tcopy(q);

  for( X = 0; X < q.N; ++X)
  {
    a = X >> m;       // top n bits
    b = X & mask;     // bottom m bits, mask = (2**m)-1 = 0111..1 (m 1's)
    c = f(a,b,z);     // perform the classical operation

    if( c > mask) error( "operate: function result overflow");

    Y = (a << m) | c; // concatenate the bits
    q.a[Y] = t.a[X];  // perform the state permutation
  }
}

// quantum FFT
//
void qfft( State q, int inv)
{
  fft( q.a, q.N, inv);  double s = 1/sqrt(q.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 < q.N; ++i)
  {
    j = bitrev( i, q.n);  

    if( i < j) { swap( q, i, j);  q.a[i] *= s;  q.a[j] *= s; }

    else if( i == j) q.a[i] *= s;
  }
}
  
// print in binary, n bits
//
void printb( unsigned long x, unsigned int n)
{
  unsigned long mask = 1LU << (n-1);

  while( mask) { putchar( (x & mask) ? '1' : '0'); mask >>= 1; }
}

// display message and qubit state or entanglement measure
//
void print( const char *msg, State q, int option)
{
  printf( "%s:%s", msg, option == PRINT_BITS ? "\n" : "");

  if( option == PRINT_TANGLE)
  {
    for( unsigned int k = 0; k < q.n; ++k) printf( " %g", tangle(q,k));
  }
  else for( unsigned long i = 0; i < q.N; ++i)
  {
    if( option == PRINT_BITS)
    {
      printf( "  "); printb( i, q.n);  // 0+x to make -0 print as just 0
      printf( " %-8.4g (%g,%g)\n", abs2(q.a[i]), 0+creal(q.a[i]), 0+cimag(q.a[i]) );
    }
    else
    {
      if( option == PRINT_LABELS) printf( " %lu:", i); else printf( " "); 

      printf( "(%g,%g)", 0+creal(q.a[i]), 0+cimag(q.a[i]) );
    }
  }
  if( option != PRINT_BITS) putchar('\n');
}

// print error message and exit
//
void error( const char *msg)
{
  fprintf( stderr, "%s\n", msg);  exit(1);
}

// magnitude squared
//
// Equivalent to abs(x)**2 but without the square-root and square operations
//
// In C++ this is available as norm(x) for complex
//
double abs2( double complex x)
{
  double r = creal(x), i = cimag(x);  return r*r + i*i;
}

// malloc with exit on error
//
void *emalloc( unsigned long size)
{
  void *r = malloc(size);  if( !r) { error( "qce: malloc() failed"); }  return r;
}