// quantum computing emulation: Shor's factoring algorithm
//
// Usage: Shor [-dE] [N [L [y]]]
//
// N = number to be factored, L = number of qubits in the work register
// 
// default L = smallest value such that 2**L >= N**2
//
// Will generate random value if y is not specified
//
//  -d  extra debug output
//  -E  show entanglement
//
// Due to randomness in the measurement step,
// results can vary for the same input parameters
//
// R. Perry, Jan. 2018
//
#include <stdio.h>
#include <stdlib.h> // atoi()
#include <tgmath.h>
#include "qce.h"

#define PMAX 8 // max number of qubits for printing unless debug specified

// y is global so f_shor() can access it when invoked from operate()
//
static unsigned long y;

// y**a mod N
//
unsigned long modpow( unsigned long y, unsigned long a, unsigned long N)
{
  unsigned long r = 1, x = y;

  while(a)
  {
    if( a & 1) r = (r*x) % N;
    x = (x*x) % N;
    a >>= 1;
  }

  return r;
}

// classical b + (y**a mod N)
//
unsigned long f_shor( unsigned long a, unsigned long b, unsigned long N)
{
  return b ^ modpow( y, a, N);
}

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

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

 // debug option
 //
  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( "Shor: bad option");
    }
    --argc, ++argv;
  }

 // number to factor
 //
  unsigned long N = 2;

  if( argc > 1) { N = atoi(argv[1]); if( N == 0) error( "Shor: bad N argument"); }

  unsigned long N2 = N*N;

  if( N2/N != N) { printf( "Shor: N**2 overflow\n"); N2 = N; }

 // L = size of work register
 //
  unsigned int L = 1;  unsigned long L2 = 2;

  while( L2 < N2) { ++L; L2 <<= 1; if( L2 == 0) error( "Shor: 2**L overflow"); }

  printf( "Shor: N = %lu, default L = %u\n", N, L);

  if( argc > 2) { L = atoi(argv[2]); if( L == 0) error( "Shor: bad L argument"); }

 // K = size of auxiliary register, 2**K >= N
 //
  unsigned int K = 1;  unsigned long K2 = 2;  while( K2 < N) { ++K; K2 <<= 1; }

 // y specified or choosen randomly 2 <= y < N-1
 //
  if( argc > 3) y = atoi(argv[3]); else y = (N > 2) ? 2 + IRAND(N-3) : 1;

  printf( "Shor: seed = %u, N = %lu, L = %u, K = %u, y = %lu\n", seed, N, L, K, y);

  State q = Q(L+K);

#if 0 // testing modpow()
  for( unsigned long a = 0; a < q.N; ++a)
    printf( "%lu %lu\n", a, modpow( y, a, N) );
  return 0;
#endif

  init( q, 0, 0); // initial state |0...0>

  if(debug) { print( "q", q, 0);  if( Eflag) print( "E", q, -1); }

 // perform H on top L qubits
 //
  for( unsigned int i = q.n-1; i >= K; --i)
  {
    H( q, i);  if(debug) { print( "H", q, 0);  if( Eflag) print( "E", q, -1); }
  }

  operate( q, L, f_shor, N); 

  if(debug) { print( "Shor", q, 0);  if( Eflag) print( "E", q, -1); }
  
 // measure and note lower K bits
 //
  unsigned long m = M(q), Kmask = ((1LU<<K)-1), k = m & Kmask;

  printf( "Shor: m = %lu, k = %lu\n", m, k);

 // could just do the FFT on q without measuring and reducing,
 // but reduction makes the state size smaller so the FFT on r is easier
 //
  State r = reduce( q, L, k);

  if( debug || r.n <= PMAX) { print( "reduce", r, 0);  if( Eflag) print( "E", r, -1); }
  else printf( "reduce\n");

  qfft( r, 0); 
  if( debug || r.n <= PMAX) { print( "FFT", r, 0);  if( Eflag) print( "E", r, -1); } 
  else printf( "FFT\n");

 // for plotting the probabilities
 //
  for( unsigned long i = 0; i < r.N; ++i) printf( "%lu %g\n", i, abs2(r.a[i]) );

 // find peaks
 //
  printf( "Peaks:\n");  double d1, d2 = 0, d3 = 0;

  for( unsigned long i = 0; i < r.N; ++i)
  {
    d1 = d2; d2 = d3; d3 = abs2(r.a[i]);
    if( d1 < d2 && d2 > d3) printf( "%lu %g\n", i-1, d2);
  }

  return 0;
}