// quantum computing emulation: Deutsch-Jozsa coin analogy
//
// Determine P(heads) for an NS-sided coin
//
// Usage: DJ-coin [-dE] [NS [NH]]
//
//  NS = number of sides
//  NH = number of heads
//  -d  extra debug output
//  -E  show entanglement
//
// R. Perry, July 2018; updated June 2020 for qce++
//
#include <iostream>
#include <cstdlib> // atol(), rand()
#include <cmath>
#include <cstring> // memset()
#include "qce.h"

using namespace qce; using namespace std;

#define Real double
#define Complex Real
#define State state<Complex,Real>

// print X probabilities, without auxiliary bit
//
void print_X( const char *msg, const State &q)
{
  cout << msg << ":\n";

  for( unsigned long i = 0; i < q.N; i+=2)
  {
    cout << "  "; printb( i>>1, q.n-1);
    cout << "  " << abs2(q.a[i])+abs2(q.a[i+1]) << "\n";
  }
}

// set state = H|0...0>|1>
//
void init_q( State &q)
{
  q.init(1); for( unsigned int k = 0; k < q.n; ++k) q.H(k);
}

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

    while( n > 1) { j = irand(n); --n; t = R[n]; R[n] = R[j]; R[j] = t; }
}

unsigned char *R; // for random and general (non-constant/non-balanced) functions
//
void print_R( unsigned long M)
{
  cout << "R: ";
  for( unsigned long i = 0; i < M; ++i) cout << (R[i] ? '1' : '0');
  cout << '\n';
}
//
// set R with NH 1's in first NS positions randomly
//
void init_R( unsigned long M, unsigned long NS, unsigned long NH)
{
  if( !R) R = (unsigned char *) emalloc(M);

  memset( R, 0, M);

  for( unsigned long i = 0; i < NH; ++i) R[i] = 1;

  shuffle( R, NS);
}

// f: the hidden Boolean function
//
// state (x,y) -> (x,y+f(x))
//
void f( State &q)
{
  for( unsigned long i = 0; i < q.N; i += 2) if( R[i>>1]) q.swap( i, i|1);
}

// debug: display state
//
void disp( const char *msg, const State &q, int Eflag)
{
  q.print( msg, PRINT_BITS); if( Eflag) q.print( "E", PRINT_TANGLE);
}

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

int main( int argc, char *argv[])
{
  int debug = 0, Eflag = 0; unsigned int seed = SRAND();

 // options
 //
  if( argc > 1 && argv[1][0] == '-') 
  {
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i])
    {
      case 'd': ++debug; break;
      case 'E': ++Eflag; break;
      default: error( "DJ-coin: bad option");
    }
    --argc, ++argv;
  }

 // number of sides and number of Heads
 //
  unsigned long NS = 2; if( argc > 1) NS = atol(argv[1]); 

  if( NS < 1) error( "DJ-coin: bad NS arg");

  unsigned long NH = 1; if( argc > 2) NH = atol(argv[2]);

  if( NH > NS) error( "DJ-coin: bad NH arg");

 // number of qubits is smallest q.n such that M = 2**(q.n-1) >= NS
 //
  unsigned int m = 1; unsigned long M = 1LU << m; while( M < NS) { ++m; M <<= 1; }

  State q(m+1);

  cout << "DJ-coin: seed = " << seed << ", q.n = " << q.n << ", q.N = " << q.N <<
    ", M = " << M << ", NS = " << NS << ", NH = " << NH << "\n";

  init_R( M, NS, NH); print_R( M);

  init_q( q); if( debug) disp( "init", q, Eflag);

 // apply f
 //
  f( q); if( debug) disp( "f", q, Eflag);

 // apply H
 //
 // can do all n bits (k=0,...) or skip the lsb (k=1,...)
 //
  for( unsigned int k = 0; k < q.n; ++k) { q.H(k); /* if( debug) disp( "H", q, Eflag); */ }

  if( debug) disp( "H", q, Eflag);

  print_X( "X", q);

  return 0;
}