// quantum computing emulation: Grover's search, analogy with Galperin colliding blocks
//
// see: "Playing Pool with |Psi>: from Bouncing Billiards to Quantum Search",
//  Adam R. Brown, 30 Oct. 2020, https://arxiv.org/abs/1912.02207
//
// Usage: Grover [-dE] [qn [m]]
//
//  qn = number of qubits
//   m = value to match
//  -d  extra debug output
//  -E  show entanglement
//
//  if m is negative or not specified then it is generated randomly
//
// R. Perry, June 2021
//
#include <iostream>
#include <cstdlib> // atoi(), atol()
#include "qce.h"

using namespace qce; using namespace std;

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

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

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

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

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

 // value to match
 //
  State q(qn); unsigned long m, M; // f(m)=1, otherwise 0

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

  if( arg < 0) m = irand(q.N); else m = arg;

  if( m >= q.N) error( "G: bad m arg");

  M = (m+1)%q.N; // any index != m, amplitudes will be equal

  cout << "G: seed = " << seed << ", q.n = " << q.n << ", q.N = " << q.N <<
    ", m = " << m << "\n";

 // initial state -H|0...0>
 //
  q.init(0); q.a[0] = -q.a[0]; for( unsigned int i = 0; i < q.n; ++i) q.H(i);

 // time, velocities, and positions for analogy with Galperin colliding blocks
 //
  Real dt, t = 0, u = q.a[m], v = q.a[M], x = 0.5, y = 1;

  unsigned int iter = 0;

  Real ztol = 10*q.eps; // tolerance for zero

  if( debug) disp( "init", q, Eflag);
   else cout << t << " " << x << " " << y << " " << abs2(u) << " " << u << " " << v << "\n";

 // collisions
 //
  while( u < 0 || u-v > ztol) // ztol to avoid finite precision problems
  {
    if( debug) cout << "iter=" << ++iter << "\n";

    if( u < 0) // m-wall collision: 0 = x + u*dt, flip sign of state m
    {
      dt = -x/u; t += dt; x = 0; y += v*dt; u = q.a[m] = -q.a[m];
      if( debug) disp( "f", q, Eflag);
    }
    else // m-M collision: x + u*dt = y + v*dt, inversion about the average
    {
      dt = (y-x)/(u-v); t += dt; x = y = x + u*dt;
      Real avg = 0; for( unsigned long i = 0; i < q.N; ++i) avg += q.a[i];
      avg *= 2.0/q.N; for( unsigned long i = 0; i < q.N; ++i) q.a[i] = avg - q.a[i];
      u = q.a[m]; v = q.a[M];
      if( debug) // check max probability
      {
        disp( "inv", q, Eflag);
        unsigned long s = 0;  Real p = abs2(q.a[0]), ai;
        for( unsigned long i = 1; i < q.N; ++i) {
          ai = abs2(q.a[i]); if( ai > p) { p = ai; s = i; } }
        if( s != m) cout << "G: s != m\n";
        cout << "P(" << s << ") = " << p << "\n";
      }
    }

    if( !debug)
      cout << t << " " << x << " " << y << " " << abs2(u) << " " << u << " " << v << "\n";
  }

  dt = 0.5; t += dt; x += u*dt; y += v*dt;

  if( !debug)
    cout << t << " " << x << " " << y << " " << abs2(u) << " " << u << " " << v << "\n";

  return 0;
}