// quantum computing emulation: counting, based on amplitude estimation
//
// Given f: X->{0,1}, count number of solutions to f(x)=1
//
// Usage: count [-d] [n [m [t [s]]]]
//
//  -d  extra debug output
//   n = number of qubits for x
//   m = number of qubits for QFT
//   t = number of solutions
//   s = start index of solutions
//
// defaults: n = m = t = 2, s = 0
//
// 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(), qsort()
#include <cmath>
#include "qce.h"

using namespace qce;
using namespace std;

#define Real double
#define Complex std::complex<Real>
#define Cstate state<Complex,Real>
#define Rstate state<Real,Real>

// comparison function for qsort: real part = index, imag part = probability
//
int compare( const void *v1, const void *v2) {
  Real p1 = ((const Complex *)v1)->imag(), p2 = ((const Complex *)v2)->imag();
  if( p1 < p2) return -1; else if( p1 > p2) return +1; else return 0;
}

// cache efficient transpose, c(N,M) <- r(M,N)
// ref: https://wgropp.cs.illinois.edu/courses/cs598-s15/lectures/lecture08.pdf
//
void transpose( Complex *c, Real *r, unsigned long M, unsigned long N, // dimensions
  unsigned long m, unsigned long n) // current partition size
{
#if 0 // testing
  cout << "transpose: " << m << " " << n << '\n';
  if( m <= 2 && n <= 2) {
#else
  if( m <= 16 && n <= 16) { // transpose the subblock
#endif
    for( unsigned long i = 0, iN = 0; i < m; ++i, iN += N)
      for( unsigned long j = 0, jM = 0; j < n; ++j, jM += M) 
        c[jM+i] = r[iN+j];
  } else if( m > n) { // subdivide the long side
    unsigned long m2 = m/2;
    transpose( c, r, M, N, m2, n); transpose( c+m2, r+m2*N, M, N, m-m2, n);
  } else {
    unsigned long n2 = n/2;
    transpose( c, r, M, N, m, n2); transpose( c+n2*M, r+n2, M, N, m, n-n2);
  } 
}

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

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

 // n,m = number of qubits for x and QFT, state = (r,x), r = 0...M, x = 0...N
 //
  unsigned int n = 2, m = 2;
  if( argc > 1) { n = atoi(argv[1]); } if( n < 1) error( "count: bad n arg");
  if( argc > 2) { m = atoi(argv[2]); } if( m < 1) error( "count: bad m arg");
  unsigned long N = 1UL << n, M = 1UL << m;
  Rstate q(n+m); // q.n = n+m, q.N = N*M

 // number of solutions and starting solution index
 //
  unsigned long t = 2, s = 0;
  if( argc > 3) { t = atol(argv[3]); } if( t >  N ) error( "count: bad t arg");
  if( argc > 4) { s = atol(argv[4]); } if( s >= N ) error( "count: bad s arg");

  cout << "count: n = " << n << ", m = " << m << ", t = " << t
    << ", s = " << s << ", q.n = " << q.n << ", q.N = " << q.N << "\n";

#if 0 // testing transpose
  for( unsigned int i = 0; i < q.N; ++i) q.a[i] = i;
  q.print( "q"); Cstate b(q.n); transpose( b.a, q.a, M, N, M, N); b.print( "b");
  return 0;
#endif

 // initial state vector: H|0>H|0> => a*[1,...,1], a = 1/sqrt(M*N)
 //
  Real a = (Real)1.0/sqrt((Real)q.N); for( unsigned long i = 0; i < q.N; ++i) q.a[i] = a;
  if( debug) q.print( "init");

 // create |j>(G**j|x>), cache-efficient qubit ordering
 //
  for( unsigned long j = 1; j < M; ++j) { 
    for( unsigned long k = j, kn = 0; k < M; ++k) { 
      kn = k << n; if( debug) cout << j << " " << k << '\n';
      for( unsigned long T = 0, S = s; T < t; ++T) { // flip sign of solution states
        unsigned long i = kn|S; q.a[i] = -q.a[i];
        S = (S+1) % N; }
      if( debug) q.print( "flip");
      Real avg = 0; for( unsigned long S = 0; S < N; ++S) avg += q.a[kn|S];
      avg *= (Real)2.0/(Real)N; for( unsigned long S = 0; S < N; ++S) {
        unsigned long i = kn|S; q.a[i] = avg - q.a[i]; }
      if( debug) q.print( "avg");
    }
  }

 // transpose to complex state for cache-efficient QFT
 //
  Cstate f(q.n); transpose( f.a, q.a, M, N, M, N); if( debug) f.print( "f");

 // inverse QFT on bottom m qubits
 //
  Complex *A = f.a; for( unsigned long S = 0; S < N; ++S, A += M) // top n qubits
    fft<Complex,Real>( A, m, 0); // inverse QFT = FFT
  if( debug) f.print( "qft");

 // compute probabilities, store in f.a[0,...,M-1]: real part = index, imag part = probability
 // (aside: C/C++ complex real,imag parts are contiguous)
 // accumulate sums in memory order of f.a[] for cache efficiency
 //
  for( unsigned long k = 0; k < M; ++k) { f.a[k].imag( abs2(f.a[k])); f.a[k].real(k); }
  for( unsigned long S = 1; S < N; ++S) for( unsigned long Sm = S << m, k = 0; k < M; ++k)
      f.a[k].imag( f.a[k].imag() + abs2(f.a[Sm|k]) );
  for( unsigned long k = 0; k < M; ++k) cout << k << " " << f.a[k].imag() << '\n';

 // sort and redisplay probabilities with count estimates round(c) ~= t
 // initial probability of match = t/N = sin(pi*k/M)**2 => exact t: c = N*sin(pi*k/M)**2
 //   =>  exact k = asin(sqrt(t/N))*M/pi, and also M-k if k > 0
 //       will be integer if t/N is 1 or 1/2, then asin(sqrt(t/N))/pi is 2 or 4
 //
  qsort( f.a, M, sizeof(Complex), compare);
  Real pt = 0; // total probability of t match
  Real mean = 0; // weighted average of unrounded c values
  unsigned long Sm = 1UL << m; // store c values in f.a[Sm|j].real
  for( unsigned long j = 0; j < M; ++j) {
    Real k = f.a[j].real(), c = sin(f.pi*k/M), p = f.a[j].imag();
    c *= c*N; mean += p*c; f.a[Sm|j].real(c); unsigned long est = round(c);
    cout << "count: " << k << " " << f.a[j].imag() << " " << c << " " << est << '\n';
    if( est == t) pt += p;
  }
  Real var = 0; for( unsigned long j = 0; j < M; ++j) {
    Real p = f.a[j].imag(), v = f.a[Sm|j].real() - mean; var += p*v*v;
  }
  Real stdev = sqrt(var), ek = asin(sqrt((Real)t/N))*M/f.pi; // exact k
  cout << "count: pt = " << pt << ", 1-pt = " << 1-pt << ", ek = " << ek << " or " << M-ek
    << ", mean = " << mean << ", stdev = " << stdev << '\n';

  return 0;
}