// quantum computing emulation: Amplitude Amplification, generalizations of Grover's algorithm
//
// Usage: AMP [-dfx] [n [k [m]]]
//
//  -d  extra debug output, -dd more debug, -ddd all debug
//  -f  use fractional final rotation, or:
//  -x  use extra qubit
//   n = problem size, n+1 qubits with -x option, otherwise n qubits
//   k = number of iterations, default = ceil(optimum#) or floor with -f
//   m = value to match
//
//  if k is negative the default is used
//  if m is negative or not specified then it is generated randomly
//
// R. Perry, June 2022
//
// Reference: Quantum Amplitude Amplification and Estimation,
// Gilles Brassard, Peter Hoyer, Michele Mosca, and Alain Tapp,
// May 2000, https://arxiv.org/abs/quant-ph/0005055
//
#include <iostream>
#include <cstdlib> // atoi(), atol()
#include <cmath>
#include "qce.h"

using namespace qce;
using namespace std;

#define Real long double
#define Complex std::complex<Real>
#define State state<Complex,Real>

// general rotation operator 2|s><s|-I, requires copy of the initial state
// same as inversion about the average if |s> = H|0...0>
// only used if GENROT is defined
//
void inv( State &q, const State &s)
{
  Complex sum = 0; for( unsigned long i = 0; i < q.N; ++i) sum += Conj(s.a[i])*q.a[i];
  sum *= 2.0; for( unsigned long i = 0; i < q.N; ++i) q.a[i] = s.a[i]*sum-q.a[i];
}

// rotation operator 2|s><s|-I for |s> with amplitudes A*[1,...,1],B*[1,...,1] 
// same as inversion about the average if N = q.N (only A used)
//
void inv( State &q, unsigned long N, Real A, Real B)
{
  Complex sumA = 0, sumB = 0;
  for( unsigned long i = 0; i < N; ++i) sumA += q.a[i];
  for( unsigned long i = N; i < q.N; ++i) sumB += q.a[i];
  Complex prod = (Real)2.0*(A*sumA + B*sumB), Aprod = A*prod, Bprod = B*prod;
  for( unsigned long i = 0; i < N; ++i) q.a[i] = Aprod-q.a[i];
  for( unsigned long i = N; i < q.N; ++i) q.a[i] = Bprod-q.a[i];
}

// fractional rotation operator
// same as inversion about the average if phi0 = pi
//
void inv( State &q, Real phi0)
{
  Complex e0 = exp(Complex(0,phi0)), sum = 0;
  for( unsigned long i = 0; i < q.N; ++i) sum += q.a[i];
  sum *= ((Real)1.0-e0)/(Real)q.N; for( unsigned long i = 0; i < q.N; ++i) q.a[i] = sum-q.a[i];
}

// display match and top two probabilities
//
void check( const State &q, unsigned long m, int iter)
{
  unsigned long z[2]; Real p[2], t; z[0]=z[1]=0; p[0]=p[1]=abs2(q.a[0]);
  for( unsigned long i = 1; i < q.N; ++i) { t = abs2(q.a[i]);
    if( t > p[0]) { p[1]=p[0]; p[0]=t; z[1]=z[0]; z[0]=i; }
    else if( t > p[1]) { p[1]=t; z[1]=i; }
  }
  Real pm = abs2(q.a[m]); cout << iter << " " << pm << " " << 1-pm << " " << m;
  for( int i = 0; i < 2; ++i)
    cout << " " << p[i] << " " << 1-p[i] << " " << z[i] << " ";
  cout << '\n';
}

// for use with fractional final rotation
//
typedef struct { Real xn, yn; unsigned long N; } Params;
//
// after k standard Grover rotations the state is x|w>+y|s'>
// xn = sqrt(N-1)*x/N, yn = y/N
//
// function to minimize to find phi0 and phi1 to get state 1|w>+0|s'>
//
Real f( const Real *z, int n, void *params)
{
  Params *p = (Params *) params;
  Real phi0 = z[0], phi1 = z[1];
  Complex e0 = exp(Complex(0,phi0)), e1 = exp(Complex(0,phi1));
  Complex r = (e1*((Real)1.0-e0))*p->xn + (-e0*(p->N-(Real)1.0)-(Real)1.0)*p->yn;
  return abs2(r);
}

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

  if( argc > 1 && argv[1][0] == '-') { // options
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i]) {
      case 'd': ++debug; break;
      case 'f': ++fflag; break;
      case 'x': ++xflag; break;
      default: error( "AMP: bad option"); }
    --argc, ++argv;
  }

  if( fflag && xflag) error( "AMP: can't do both -f and -x");

 // n = problem size, N = 2**n, n+1 qubits with -x option
 //
  unsigned int n = 2, qn; if( argc > 1) n = atoi(argv[1]); 
  if( n < 1) error( "AMP: bad n arg");
  unsigned long N = 1UL << n; if( xflag) qn = n+1; else qn = n;
  State q(qn);

 // number of iterations: r = optimal but not integer except for N=4
 //
  Real theta = 2*asin(1/sqrt(N)), r = (q.pi/2)/theta - 0.5;
  long arg = -1; if( argc > 2) arg = atol(argv[2]);
  int k; if( arg >= 0) k = arg; else if( fflag) k = floor(r); else k = ceil(r);

 // value to match
 //
  unsigned long m; arg = -1; if( argc > 3) arg = atol(argv[3]);
  if( arg < 0) m = irand(N); else m = arg;
  if( m >= N) error( "AMP: bad m arg");

  cout << "AMP: seed = " << seed << ", n = " << n << ", N = " << N <<
    ", q.n = " << q.n << ", q.N = " << q.N <<
    ", r = " << r << ", k = " << k << ", m = " << m << "\n";

  Real A, B, a = 1/sqrt(N); // a for equal amplitudes

 // initial state vector:    H|0...0> => a*[1,...,1] 
 // or with -x option:   X|0>H|0...0> => A*[1,...,1],B*[1,...,1] 
 // where X = [A B; B -A]/a
 //
  q.init(0); if( xflag) {
    theta = (q.pi/2)/(ceil(r)+0.5); A = sin(theta/2); B = sqrt((a-A)*(a+A));
    for( unsigned long i = 0; i < N; ++i) q.a[i] = A;
    for( unsigned long i = N; i < q.N; ++i) q.a[i] = B;
  } else { A = a; B = 0;
    for( unsigned long i = 0; i < N; ++i) q.a[i] = a;
  }
  if( debug) q.print( "init");
  check( q, m, 0);

#ifdef GENROT
  State s(q.n); memcpy( s.a, q.a, q.N*sizeof(Complex)); // copy of initial state
#endif
  
  int iter; for( iter = 1; iter <= k; ++iter) {
    q.a[m] = -q.a[m]; // flip sign of state m
    if( debug) q.print( "flip");
#ifdef GENROT
    inv(q,s);
#else
    inv(q,N,A,B);
#endif
    if( debug) q.print( "inv");
    check( q, m, iter);
  }

 // using fractional final rotation
 //
  if( fflag) { // after k standard Grover rotations the state is x|w>+y|s'>
    Real z[2], angle = (k+0.5)*theta, x = sin(angle), y = cos(angle);
    if( debug > 1) cout << "x = " << x << ", y = " << y << '\n';
    Params params = { (Real)(sqrt(N-1)*x/N), y/N, N };
#if 0 // grid search for approximate minimum of f()
    // debug>2 output can be used for surface plot
    int p0 = 0, p1 = 0; Real dp = 2*q.pi/10, _fmin = HUGE_VAL;
    for( int i = 0; i < 10; ++i) { z[0] = i*dp; // phi0
      for( int j = 0; j < 10; ++j) { z[1] = j*dp; // phi1
        Real fval = f(z,N,(void *)&params); if( fval < _fmin) { _fmin = fval; p0 = i; p1 = j; }
	if( debug > 2) cout << i << " " << j << " " << fval << '\n';
    }}
    if( debug > 1) cout << "p0 = " << p0 << ", p1 = " << p1 << ", _fmin = " << _fmin << '\n';
    z[0] = p0*dp; z[1] = p1*dp;
#else // just use simplex
    z[0] = z[1] = q.pi/4;
#endif
   // Nelder-Mead simplex search for phi0 and phi1 to minimize f()
    simplex<Real> S( f, z, 2, (void *)&params, q.pi/2, 10*q.eps, q.eps*q.eps, debug > 2);
    Real phi0 = S.x[S.L][0], phi1 = S.x[S.L][1], fmin = S.f[S.L];
    if( debug > 1) cout << "phi0 = " << phi0 << ", phi1 = " << phi1 << ", fmin = " << fmin << 
      ", count = " << S.count << ", f = " << f(S.x[S.L],N,(void*)&params) << '\n';
   // apply the rotation
    q.a[m] = exp(Complex(0,phi1))*q.a[m]; if( debug) q.print( "flip");
    inv( q, phi0); if( debug) q.print( "inv");
    check( q, m, iter);
  }

  return 0;
}