// quantum computing emulation: Bernstein & Vazirani problem
//
// Determine the hidden bit string (a) of a parity function:
//
//   f(x) = a*x + b
//
// Usage: BV [-ersmdE0] [qn [p [a [b [nt]]]]]
//
//  qn = number of qubits for x; qn+1 total qubits
//   p = probability of depolarizing noise, or non-resonant pulse error level
//  nt = number of trials for -s option, default = 100
//  -e  non-resonant pulse error
//  -r  use random global phase factor
//  -s  generate statistics: correct bits vs. number of trials
//  -m  generate measurement probability table
//  -d  extra debug output, -dd for even more output
//  -E  show entanglement
//  -0  initialize result qubit to |0> instead of H|1> = |->
//
//  if a and/or b are negative or not specified then they are generated randomly
//
// R. Perry, April 2018
//
#include <stdio.h>
#include <stdlib.h> // atoi()
#include <tgmath.h>
#include <string.h> // memset()
#include "qce.h"

#if 0
// temporary debug - sum of amplitudes magnitude squared
//
double sum( State q)
{
  double R = 0;

  for( unsigned long i = 0; i < q.N; ++i) R += abs2(q.a[i]); 

  return R;
} 
#endif

double pi;

// classical a*x + b
//
unsigned long f( unsigned long x, unsigned long a, unsigned long b)
{
  unsigned long r = b;

  x &= a;  while( x) { r ^= (x & 1); x >>= 1; }

  return r;
}

#if 0
// b would have to be global so fxy() can access it when invoked from operate()
//
// unsigned long b;

// classical y + a*x + b
//
unsigned long fxy( unsigned long x, unsigned long y, unsigned long a)
{
  return y ^ f(x,a,b);
}
#endif

// depolarize all n qubits - just before measuring
//
// depolarization produces a mixed state which can not be represented using a state vector,
// so we fudge it here and just create a state which will lead to proper measurements
//
// each state amplitude a is changed using a convex combination:
//   magnitude -> sqrt( (1-p)*|a|**2 + p*|h|**2 )
//   phase -> (1-p)*arg(a) + p*0
// where h = 1/sqrt(2**n)
//
void depolarize( State q, double p)
{
  double mag2, ang, newang, h2 = 1 / (double) q.N;

  // printf( "h2 = %g\n", h2);

  for( unsigned long i = 0; i < q.N; ++i)
  {
    mag2 = abs2(q.a[i]);  ang = carg(q.a[i]);  newang = (1-p)*ang;
  
    // printf( "mag2 = %g, ang = %g, newang = %g\n", mag2, ang, newang);

    q.a[i] = sqrt( (1-p)*mag2 + p*h2) * CMPLX( cos(newang), sin(newang) );
  }
}

//---------------------------------------------------------------------

int main( int argc, char *argv[])
{
  int rphase = 0, stats = 0, eflag = 0, mflag = 0, Eflag = 0, r = 1, debug = 0;

  pi = 4*atan(1);  unsigned int seed = SRAND();

 // command-line options
 //
  if( argc > 1 && argv[1][0] == '-') 
  {
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i])
    {
      case 'e': ++eflag; break;
      case 'r': ++rphase; break;
      case 's': ++stats; break;
      case 'm': ++mflag; break;
      case 'd': ++debug; break;
      case 'E': ++Eflag; break;
      case '0': r = 0; break;
      default: error( "BV: bad option");
    }
    --argc, ++argv;
  }

 // number of qubits for x
 //
  unsigned int qn = 2;  if( argc > 1) qn = atoi(argv[1]);

  if( qn < 1) error( "BV: bad qn arg");

  unsigned long N = 1LU<<qn;  State q = Q(qn+1);

 // probability of depolarizing noise, or non-resonant pulse error level
 //
  double p = 0, ra = 0;  if( argc > 2) ra = atof(argv[2]); // ra can be negative

  if( !eflag) { p = ra; ra = 0; }

  if( p < 0 || p > 1) error( "BV: bad p arg");

 // parity function parameters
 //
  unsigned long a, b;

  long arg = -1;  if( argc > 3) arg = atol(argv[3]);

  if( arg < 0) a = 1+IRAND(N-1); else a = arg;  if( a >= N) error( "BV: bad a arg");

  arg = -1;  if( argc > 4) arg = atol(argv[4]);

  if( arg < 0) b = IRAND(2); else b = arg;  if( b >= 2) error( "BV: bad b arg");

 // number of trials for stats option
 //
  unsigned int nt = 100;  if( argc > 5) nt = atoi(argv[5]);

 // final measured value should equal (a,1)
 //
  unsigned long m = (a << 1) | 1;

  printf( "BV: seed = %u, qn = %u, N = %lu, p = %g, ra = %g, a = %lu, b = %lu, m = %lu\n",
    seed, qn, N, p, ra, a, b, m);

 // initial state |0...0>|r>
 //
  init( q, r, rphase);  if( debug > 1) { print( "q", q, 0); if( Eflag) print( "E", q, -1); }

 // fprintf( stderr, "sum = %g\n", sum(q) ); // temporary debug

 // optionally apply H to the result register
 //
  if(r) { Hra( q, 0, ra);  if( debug > 1) {
      print( "Hr", q, 0); if( Eflag) print( "E", q, -1); } }

 // fprintf( stderr, "sum = %g\n", sum(q) ); // temporary debug

 // apply H to the control register
 //
  for( unsigned int k = 1; k < q.n; ++k) { Hra( q, k, ra);
    if( debug > 1) { print( "Hc", q, 0); if( Eflag) print( "E", q, -1); } }

 // fprintf( stderr, "sum = %g\n", sum(q) ); // temporary debug

 // perform (x,y) -> (x,y+f(x))
 //
 // Could use: operate( q, qn, fxy, a); // with b global
 // which makes a copy of the state and can perform a general permutation.
 //
 // But applying f directly is simpler:
 //
  for( unsigned long i = 0; i < q.N; i += 2) if( f(i>>1,a,b) ) swap( q, i, i|1);

  if( debug > 1) { print( "f", q, 0); if( Eflag) print( "E", q, -1); }

 // depolarizing noise should go here, but it is not affected by H
 // (or any unitary transformation U, since U operating on a depolarized
 // density matrix I is U*I*U' = I), so do it after H

 // apply H to all of the qubits
 //
  for( unsigned int k = 0; k < q.n; ++k) { Hra( q, k, ra);
    if( debug > 1) { print( "H", q, 0); if( Eflag) print( "E", q, -1); } }

 // depolarizing noise
 //
  if( p > 0) depolarize( q, p);

 // stats option
 //
 // Probability of exact correct measurement is not interesting
 // since it is just a linear function of noise probability
 // going from 1.0 to 0 (or 0.5 to 0 with -0 option) as p goes from 0 to 1.
 //
 // So instead we do a Monte-Carlo simulation, following
 // "Quantum learning robust against noise", Andrew W. Cross, Graeme Smith,
 // and John A. Smolin, Phys. Rev. A 92, 012327, July 2015. 
 //
  if( stats)
  {
    unsigned long z, e;		// raw measurement value and estimate of parameter a
    unsigned int *t, ones[qn];	// counter for number of 1's in each bit position

    memset( ones, 0, qn*sizeof(int));

    for( unsigned int trials = 1; trials <= nt; ++trials)
    {
      do { z = M(q); } while( (z & 1) == 0); // skip if result register is 0

      e = z >> 1;  if( debug) printf( "z = %lu, e = %lu\n", z, e);

      t = ones;  while( e > 0) { if( e & 1) { ++*t; }  e >>= 1;  ++t; } 

      if( debug) { for( unsigned int i = qn; i > 0; --i) { printf( " %u", ones[i-1]); } putchar('\n'); }

      unsigned long E = 0; // bit-wise majority estimate of parameter a

      unsigned int t2 = trials>>1;

      for( unsigned int i = qn; i > 0; --i) { E <<= 1; if( ones[i-1] > t2) E |= 1; }

      if( debug) printf( "E = %lu\n", E);

      E ^= a;  unsigned int match = qn; // check for mismatched bits

      while( E) { if( E & 1) { --match; } E >>= 1; }

      printf( "%u %u\n", trials, match);
    }

    return 0;
  }

  if( debug) { print( "q", q, 0); if( Eflag) print( "E", q, -1); }

  printf( "BV: q[%lu] = (%g,%g), abs2 = %g\n", m, creal(q.a[m]), cimag(q.a[m]), abs2(q.a[m]) );

// the measurement probabilities should be:
//
//        P(m) = (1-p)*1.0 + p/q.N, all other P(i) = p/q.N
//
// or with -0 option:
//
// P(0) = P(m) = (1-p)*0.5 + p/q.N, all other P(i) = p/q.N

  if( mflag) // print measurement probabilities
  {
    for( unsigned long i = 0; i < q.N; ++i) printf( "%lu %g\n", i, abs2(q.a[i]) );
  }

  return 0;
}