// quantum computing emulation: addition
//
// q = (x,y) -> (x,(x+y)%M)
//
// n qubits, n = 2*m, x = top m bits, y = bottom m bits, M = 2**m
//
// Usage: add [-dnpuCESab] [m [s]]
//
//   m  total number of qubits is 2*m
//   s  first partition size for SVD, min = 1, max = 2*m-1, default = m
//  -d  extra debug output
//  -n  normalize
//  -p  product state with mostly distinct amplitudes or
//  -u   uniform superposition
//  -C  do CX(0,1) if uniform
//  -E  show entanglement
//  -S  do SVD before and after add, -SS show all SVD results, -SSS check unitary
//  -a  with -S, just do SVD after add
//  -b  with -S, just do SVD before add
//
// Run with options -pE to show that addition creates entangelment
//
// -S cancels most other output so results can be run by octave.
// To reconstruct A: AA = Q(1:r,:)'*diag(S,r,r)*V(1:r,:)
//
// R. Perry, January 2022
//
#include <iostream>
#include <cstdlib> // atoi()
#include "qce.h"

using namespace qce; using namespace std;

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

//---------------------------------------------------------------------
//
int main( int argc, char *argv[])
{
  int debug = 0, nflag = 0, pflag = 0, uflag = 0, Cflag = 0, Eflag = 0, Sflag = 0;

  int aflag = 1, bflag = 1;

  if( argc > 1 && argv[1][0] == '-') { // options
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i]) {
      case 'd': ++debug; break;
      case 'n': ++nflag; break;
      case 'p': ++pflag; break;
      case 'u': ++uflag; break;
      case 'C': ++Cflag; break;
      case 'E': ++Eflag; break;
      case 'S': ++Sflag; break;
      case 'a': aflag = 1; bflag = 0; break;
      case 'b': aflag = 0; bflag = 1; break;
      default: error( "add: bad option");
    }
    --argc, ++argv;
  }

  int check = !nflag && !pflag && !uflag && !Cflag; // check addition result consistency
 //
 // number of qubits = 2*m
 //
  unsigned int m = 2; if( argc > 1) m = atoi(argv[1]);
  if( m < 1) error( "add: bad m arg");

  unsigned int s = m; if( argc > 2) s = atoi(argv[2]);
  if( s < 1 || s > 2*m-1) error( "add: bad s arg");

  State q(2*m); unsigned long M = 1LU << m, mask = M-1; // mask = 0111..1 (m 1's)

  if( !Sflag)
    cout << "add: q.n = " << q.n << ", m = " << m << ", M = " << M << ", s = " << s << "\n";

  if( uflag ) { // uniform superposition
    q.init( 0); q.H( 0); if( Cflag) q.CX( 0, 1);
    for( unsigned int k = 1; k < q.n; ++k) q.H(k);
    if( !Sflag) q.print( "before", PRINT_BITS);
  } else if( pflag) { // product state with mostly distinct amplitudes
    for( unsigned long i = 0; i < q.N; ++i) {
      q.a[i] = 1; for( unsigned int k = 0; k < q.n; ++k)
	if( i & (1UL << k)) q.a[i] *= 2*k + 2; else q.a[i] *= 2*k + 1; }
    if( !Sflag) q.print( "before", PRINT_BITS);
  } else { // set q.a[i] based on the sum so we can easily check the results
    if( !Sflag) cout << "before, a[i] = sum:\n";
    for( unsigned long i = 0; i < q.N; ++i) {
      unsigned long x = i >> m, y = i & mask, sum = (x + y) & mask, k = (x << m) | sum;
      q.a[i] = k; // unnormalized
      if( !Sflag) {
	cout << " "; printb( x, m); cout << " + "; printb( y, m); cout << " -> ";
	printb( x, m); cout << " "; printb( sum, m); cout << " : " << k << "\n"; }
      }
  }

  if( nflag) { q.normalize(); if( !Sflag) q.print( "normalized", PRINT_BITS); }

  if( Eflag) q.print( Sflag ? "# E" : "E", PRINT_TANGLE);

  if( Sflag && bflag) {
   // copy, don't overwrite q.a; partition (s,2*m-s) qubits:
    SVD a( q.a, 1UL<<s, 1UL<<(2*m-s), 1);
    cout << "# before:\n";
    if( Sflag > 1) { matprint( cout, "A", q.a, a.n, a.m); cout << a; }
     else matprint( cout, "S", a.S, 1, a.n);
    if( Sflag > 2) {
      a.check(); matprint( cout, "AA", a.AA, a.n, a.m);
      cerr << "eps = " << a.eps << '\n';
      cerr << "|A*A'-I| = " << ucheck( q.a, a.n, a.m, a.eps) << '\n';
      cerr << "|Q*Q'-I| = " << ucheck( a.Q, a.n, a.n, a.eps) << '\n';
      cerr << "|V*V'-I| = " << ucheck( a.V, a.rank, a.m, a.eps) << '\n';
      cerr << "|A-Q'*S*V| = " << matdiff( q.a, a.AA, a.n, a.m, a.eps) << '\n';
    }
  }

 //
 // perform the addition
 //
  for( unsigned long x = 1; x < M; ++x) { // skip x=0
    unsigned long start = 0, i = start, j = (i+x) & mask, count = M, xm = x << m;
    Complex ti = q.a[xm], tj;
  /*
     circular shift by x, a[i] -> a[j]
     M = 8, x = 1, j = 1, 2, 3, 4, 5, 6, 7, 0
     M = 8, x = 2, j = 2, 4, 6, 0, 3, 5, 7, 1
     M = 8, x = 4, j = 4, 0, 5, 1, 6, 2, 7, 3
  */
    if( debug) cerr << "x = " << x << "\n";
    while( 1) { // until count=0
      if( debug) cerr << i << " " << j << "\n";
      unsigned long xj = xm | j;
      tj = q.a[xj]; q.a[xj] = ti; ti = tj; --count; if( count == 0) break;
      if( j != start) i = j; else { ++start; i = start; ti = q.a[xm|i];
       if( debug) cerr << "start = " << start << "\n"; }
      j = (i+x) & mask; }
  }
 //
 // display results
 //
  if( !Sflag) {
    if( check) cout << "after, a[i] = i:\n"; else cout << "after:\n";
    for( unsigned long i = 0; i < q.N; ++i) {
      unsigned long x = i >> m, y = i & mask;
      cout << " "; printb( x, m); cout << " "; printb( y, m); cout << " : " << q.a[i] << "\n";
      if( check && q.a[i] != i) cout << "error: a[i] != i\n";
    }
  }
  if( Eflag) q.print( Sflag ? "# E" : "E", PRINT_TANGLE);

  if( Sflag && aflag) {
   // copy, don't overwrite q.a; partition (s,2*m-s) qubits:
    SVD a( q.a, 1UL<<s, 1UL<<(2*m-s), 1);
    cout << "# after:\n";
    if( Sflag > 1) { matprint( cout, "A", q.a, a.n, a.m); cout << a; }
     else matprint( cout, "S", a.S, 1, a.n);
    if( Sflag > 2) {
      a.check(); matprint( cout, "AA", a.AA, a.n, a.m);
      cerr << "eps = " << a.eps << '\n';
      cerr << "|A*A'-I| = " << ucheck( q.a, a.n, a.m, a.eps) << '\n';
      cerr << "|Q*Q'-I| = " << ucheck( a.Q, a.n, a.n, a.eps) << '\n';
      cerr << "|V*V'-I| = " << ucheck( a.V, a.rank, a.m, a.eps) << '\n';
      cerr << "|A-Q'*S*V| = " << matdiff( q.a, a.AA, a.n, a.m, a.eps) << '\n';
    }
  }

  return 0;
}