// function templates for complex or real state amplitudes
// uses macros QTYPE, qtype, and State, defined in qce.h
//
// this is included twice by qce.c, once for QTYPE Complex and once for QTYPE Real
//
// R. Perry, Aug. 2019

//------------------------------ private ------------------------------
//
//  t		temporary state copy used in operate()
//  tcopy	make temporary state copy
//
#define _T_ _name_( t, QTYPE)
//
static State _T_;
//
static void _name_( tcopy, QTYPE)( State q)
{
  if( !_T_.a) _T_ = _name_( Q, QTYPE)(q.n);
   else if( _T_.n != q.n) { free(_T_.a); _T_ = _name_( Q, QTYPE)(q.n); }

  memcpy( _T_.a, q.a, q.N*sizeof(qtype));
}

//------------------------------ public ------------------------------
//
// allocate a new state vector, uninitialized
//
State _name_( Q, QTYPE)(unsigned 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(qtype);

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

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

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

void _name_( init, QTYPE)( State q, unsigned long k
#ifdef QTYPE_COMPLEX
 , int rphase
#endif
)
{
  qtype v = 1;

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

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

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

// tensor product
//
State _name_( product, QTYPE)( State r, State s)
{
  State q = _name_( Q, QTYPE)(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 _name_( tangle, QTYPE)( State q, unsigned int k)
{
  qtype 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;

#ifdef QTYPE_COMPLEX
    prod += conj(q.a[i])*q.a[j];
#else
    prod += q.a[i]*q.a[j];
#endif

    A += _name_( abs2, QTYPE)(q.a[i]);  B += _name_( abs2, QTYPE)(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 - _name_( abs2, QTYPE)(prod)/(A*B);

  return r;
}

// Hadamard transformation on qubit k
//
void _name_( H, QTYPE)( State q, unsigned int k)
{
  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;  qtype 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)
{
  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) {
      qtype t0 = q.a[j], t1 = q.a[jk2];
      q.a[j] = s2*(t0+t1);  q.a[jk2] = s2*(t0-t1);
    }
  }
}
#endif

// NOT transformation on qubit k
//
void _name_( X, QTYPE)( 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;  _name_( swap, QTYPE)( q, i, j);

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

// Z transformation on qubit k
//
void _name_( Z, QTYPE)( 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 _name_( CX, QTYPE)( 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

    _name_( swap, QTYPE)( 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 _name_( M, QTYPE)( State q)
{
  double P = DRAND, R = 0; // R will be cumulative probability

  for( unsigned long i = 0; i < q.N; ++i)
  {
    R += _name_( abs2, QTYPE)(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 _name_( reduce, QTYPE)( State q, unsigned int n, unsigned long m)
{
  if( n >= q.n) error( "reduce: n >= q.n");

  State r = _name_( Q, QTYPE)(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 += _name_( abs2, QTYPE)(q.a[j]);

#if 0 // debug
  #ifdef QTYPE_COMPLEX
    printf( "reduce: i = %lu, j = %lu, q.a[j] = (%g,%g)\n", i, j, creal(q.a[j]), cimag(q.a[j]) );
  #else
    printf( "reduce: i = %lu, j = %lu, q.a[j] = (%g)\n", i, j, q.a[j] );
  #endif
#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 _name_( collapse, QTYPE)( 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 += _name_( abs2, QTYPE)(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 _name_( operate, QTYPE)( 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");

  _name_( tcopy, QTYPE)(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
  }
}

// display message and qubit state or entanglement measure
//
void _name_( print, QTYPE)( 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", _name_( tangle, QTYPE)(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
#ifdef QTYPE_COMPLEX
      printf( " %-8.4g (%g,%g)\n", abs2_Complex(q.a[i]), 0+creal(q.a[i]), 0+cimag(q.a[i]) );
#else
      printf( " %-8.4g (%g)\n", abs2_Real(q.a[i]), 0+q.a[i] );
#endif
    }
    else
    {
      if( option == PRINT_LABELS) printf( " %lu:", i); else printf( " "); 
#ifdef QTYPE_COMPLEX
      printf( "(%g,%g)", 0+creal(q.a[i]), 0+cimag(q.a[i]) );
#else
      printf( "(%g)", 0+q.a[i] );
#endif
    }
  }
  if( option != PRINT_BITS) putchar('\n');
}

// 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 _name_( abs2, QTYPE)( qtype x)
{
#ifdef QTYPE_COMPLEX
  double r = creal(x), i = cimag(x);  return r*r + i*i;
#else
  return x*x;
#endif
}