// quantum computing emulation: Deutsch-Jozsa problem
//
// Determine if a hidden Boolean function is constant or balanced
//
// Usage: DJ [-ersE] [qn [rp]]
//
//  qn = number of qubits
//  rp = phase shift relative value, 1 = no error
//  -e  phase shift error statistics
//  -r  use random global phase factor
//  -s  generate statistics: P0 vs. number of 1's
//  -E  show entanglement
//
// R. Perry, March 2018
//
#include <stdio.h>
#include <stdlib.h> // rand()
#include <tgmath.h>
#include <string.h> // memset()
#include "qce.h"

#define T 7 // number of test cases

double pi;

#if 0
//
// initialize state to H|00...01>
// unused, replaced by: for( unsigned int k = 0; k < n; ++k) H( q, k);
//
void H01_init( State q)
{
  double scale = 1/sqrt(q.N);

  for( unsigned long i = 0; i < q.N; ++i) { q.a[i] = scale; scale = -scale; }
}
#endif

// put array in random order
//
void shuffle( unsigned char R[], unsigned long M)
{
    unsigned char t;  unsigned long j;

    while( M > 1) { j = IRAND(M); --M; t = R[M]; R[M] = R[j]; R[j] = t; }
}

unsigned char *R; // for random and general (non-constant/non-balanced) functions
//
void print_R( unsigned long M)
{
#if 0 // debug
  printf( "  R: ");
  for( unsigned long i = 0; i < M; ++i) putchar( R[i] ? '1' : '0');
  putchar('\n');
#endif
}
//
// set R with count 1's (in position pos if random == 0)
//
void init_R_ones( State q, unsigned long count, unsigned long pos, int random)
{
  unsigned int m = q.n-1;  unsigned long M = 1LU<<m;

  if( !R) R = emalloc(M);

  memset( R, 0, M);

  if( random) { for( unsigned long i = 0; i < count; ++i) { R[i] = 1; }  shuffle( R, M); }

  else R[pos] = 1;

  print_R(M);
}
//
// set R random/balanced
//
void init_R_balanced( State q)
{
  unsigned int m = q.n-1;  unsigned long M = 1LU<<m, M2 = M>>1;

  if( !R) R = emalloc(M);

 // first half of R = 0, second half = 1
 //
  memset( R, 0, M2);  memset( R+M2, 1, M2);

 // put R in random order
 //
  shuffle( R, M);

  print_R(M);
}


// parity: XOR of the bits
//
int parity( unsigned long x)
{
  unsigned int r = 0;  while( x) { r ^= (x & 1); x >>= 1; }

  return r;
}

// f: the hidden Boolean function
//
// state (x,y) -> (x,y+f(x))
//
void f( int j, State q)
{
  switch( j) // 0 ... (T-1) for T test cases; and T for stats option
  {
   //
   // constant, only two cases
   //
    case 0: // f(x) = 0, do nothing
	break;
    case 1: // f(x) = 1, XOR last bit, i.e. swap adjacent
	for( unsigned long i = 0; i < q.N; i += 2) swap( q, i, i|1);
	break;
   //
   // balanced, four examples
   //
    case 2: // f(x) = (x >= N/2)
	for( unsigned long i = q.N >> 1; i < q.N; i += 2) swap( q, i, i|1);
	break;
    case 3: // f(x) = (x < N/2)
	for( unsigned long i = 0, mid = q.N >> 1; i < mid; i += 2) swap( q, i, i|1);
	break;
    case 4: // f(x) = parity(x)
	for( unsigned long i = 0; i < q.N; i += 2) if( parity(i>>1)) swap( q, i, i|1);
	break;
    case 5: // f(x) = random balanced function
	init_R_balanced( q);
	for( unsigned long i = 0; i < q.N; i += 2) if( R[i>>1]) swap( q, i, i|1);
	break;
   //
   // non-constant, non-balanced
   //
    case 6: // f(x) = (x < N/4)
	for( unsigned long i = 0, lim = q.N >> 2; i < lim; i += 2) swap( q, i, i|1);
	break;
   //
   // for stats option: general non-constant, non-balanced; call init_R_ones() first
   //
    case T:
	for( unsigned long i = 0; i < q.N; i += 2) if( R[i>>1]) swap( q, i, i|1);
	break;
  }
}

// P0: probability that the first n-1 qubits are zero, i.e. state |00...0x>
//
// P0=1 indicates that the hidden function is constant
//
double P0( State q)
{
  return abs2(q.a[0]) + abs2(q.a[1]);
}

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

int main( int argc, char *argv[])
{
  int rphase = 0, stats = 0, eflag = 0, Eflag = 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 'E': ++Eflag; break;
      default: error( "DJ: bad option");
    }
    --argc, ++argv;
  }

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

  if( qn < 2) error( "DJ: bad qn arg");

 // phase shift relative value
 //
  double rp = 1;  if( argc > 2) rp = atof(argv[2]);

  if( rp < 0 || rp > 2) error( "DJ: bad rp arg");

  State q = Q(qn);

  printf( "DJ: seed = %u, q.n = %u, q.N = %lu, rp = %g\n", seed, q.n, q.N, rp);

 // generate phase shift error statistics using balanced parity function (case 4)
 //
  if( eflag)
  {
    for( int i = 0; i < 121; ++i)
    {
      rp = 0.4 + i*0.01;

      init( q, 1, rphase);  for( unsigned int k = 0; k < q.n; ++k) Hrp( q, k, rp);

      f( 4, q);  for( unsigned int k = 0; k < q.n; ++k) Hrp( q, k, rp);

      printf( "%g %g\n", rp, P0(q));
    }

    return 0;
  }

 // generate statistics: P0 vs. number of 1's
 //
  if( stats)
  {
    unsigned int m = q.n-1;  unsigned long M = 1LU<<m;

   // init_R_ones( q, 1, i, 0); // one 1, position i
   // results show P0 = ((2**(m-1)-1)/(2**(m-1)))**2 independent of the position of the single 1
   //
    for( unsigned long i = 0; i <= M; ++i)
    {
      init_R_ones( q, i, 0, 1); // i 1's, random positions

      init( q, 1, rphase);  for( unsigned int k = 0; k < q.n; ++k) Hrp( q, k, rp);

      f( T, q);  for( unsigned int k = 0; k < q.n; ++k) Hrp( q, k, rp);

      // print( "q", q, 0); // tmp debug

      printf( "%lu %g\n", i, P0(q));
    }

    return 0;
  }

 // test T cases
 //
  for( int i = 0; i < T; ++i)
  {
    printf( "i=%d\n", i);

    init( q, 1, rphase);  print( "q", q, 0); // initial state |0...0>|1>

    if( Eflag) print( "E", q, -1);

    for( unsigned int k = 0; k < q.n; ++k)
    {
      Hrp( q, k, rp);  print( "H", q, 0);  if( Eflag) print( "E", q, -1);
    }

    f( i, q);  print( "f", q, 0);  if( Eflag) print( "E", q, -1);

   // can do all n bits, and verify phase in Abbott paper bottom page 3
   //
    for( unsigned int k = 0; k < q.n; ++k)
    {
      Hrp( q, k, rp);  print( "H", q, 0);  if( Eflag) print( "E", q, -1);
    }

   // or can skip the lsb, as in OSP 2018/A2P2:
   //
   // for( unsigned int k = 1; k < q.n; ++k)
   // {
   //    Hrp( q, k, rp);  print( "H", q, 0);  if( Eflag) print( "E", q, -1);
   // }

   // if P0 = 1 then f is constant, otherwise f is balanced
   //
    double p0 = P0(q);

    printf( "%3s: %g, 1-P0: %g\n", "P0", p0, 1-p0);
  }

  return 0;
}