/*****************************************************************************
 *
 * Grover's Search Algorithm
 *
 * Implements Grover's Search algorithm in the Scaffold programming language
 *
 * February 2013
 *
 *****************************************************************************/

#include <math.h>

#define n 4 // problem size (log of database size)

#define pi 3.141592653589793238462643383279502884197

// scaff_module prototypes
void diffuse(qbit q[n], qbit x[n-1]);
void Sqr(qbit a[n], qbit b[n]);
void EQxMark(qbit b[n], qbit x[n-1]);

scaff_module cRz ( qbit ctrl, qbit target, const double angle ) {
    /* cRz identity matrix:
      [ [ 1 0 0     			0     ]
        [ 0 1 0     			0     ]
        [ 0 0 e^(-i*angle/2)   0     ]
        [ 0 0 0 				e^(i*angle/2) ] ]
    */

    S(ctrl);
    CNOT(ctrl, target);
    Rz(target,-angle/2);
    CNOT(ctrl, target);
    Rz(target,angle/2);
}

/***************************
 Diffusion scaff_module
***************************/
void diffuse(qbit q[n], qbit x[n-1]) {
  int j;
  // local qbit x[n-1];  // scratch register
  // No local registers yet
  // qbit x[n-1];  // scratch register

  // Hadamard applied to q
  // No forall yet
  // forall(j=0; j<n; j++) { H(q[j]); }
  for(j=0; j<n; j++) { H(q[j]); }

  // Want to phase flip on q = 00...0
  // So invert q to compute q[0] and q[1] and ... and q[n-1]
  for(j = 0; j < n; j++)
    X(q[j]);

  // Compute x[n-2] = q[0] and q[1] and ... and q[n-1]
  // No forall yet
  // forall(j=0; j<n-1; j++) PrepZ(x[j]);
  // for(j=0; j<n-1; j++) PrepZ(x[j],0);

  Toffoli(q[1], q[0], x[0]);
  for(j = 1; j < n-1; j++)
    Toffoli(x[j-1], q[j+1], x[j]);

  // Phase flip conditioned on x[n-2]
  Z(x[n-2]);  // Phase Flip==Z if q=00...0, i.e. x[n-2]==1
  // cRz (x[n-2], q[n-1], pi);  // Phase Flip==Z if q=00...0, i.e. x[n-2]==1

  // Undo the local registers
  for(j = n-2; j > 0; j--)
    Toffoli(x[j-1], q[j+1], x[j]);
  Toffoli(q[1], q[0], x[0]);

  // Restore q
  for(j = 0; j < n; j++)
    X(q[j]);

  // Complete the diffusion
  // No forall yet
  // forall(j=0; j<n; j++) { H(q[j]); }
  for(j=0; j<n; j++) { H(q[j]); }
}

/***********************************************
 Test if the input polynomial
   b(x) = b0 + b1*x + ... + b_(n-1)*x^(n-1) == x
 over the ring GF(2)[x] mod (x^n + x + 1).


 if(tF!=0) set return result in qubit t[0] else Z
************************************************/
void EQxMark(qbit b[n], qbit x[n-1]) {
  int j;
  // No local registers yet
  // local qbit x[n-1];  // scratch register
  // qbit x[n-1];  // scratch register

  // Change b to reflect testing for the polynomial x
  for(j = 0; j < n; j++)
    if(j!=1) X(b[j]);

  // Compute x[n-2] = b[0] and b[1] and ... and b[n-1]
  // No forall yet
  // forall(j=0; j<n-1; j++) PrepZ(x[j]);
  // for(j=0; j<n-1; j++) PrepZ(x[j],0);
  Toffoli(b[1], b[0], x[0]);
  for(j = 1; j < n-1; j++)
    Toffoli(x[j-1], b[j+1], x[j]);

  // Either return result in t or Phase flip conditioned on x[n-2]
  // if(tF!=0) {
  //   CNOT(x[n-2], t[0]); // Returns result in t
  // }
  // else {
  //   // CNOT(x[n-2], t[0]); // Returns result in t
  //   Z(x[n-2]);  // Phase Flip==Z if b=01...0 == 'x', i.e. x[n-2]==1
  // }
  Z(x[n-2]);  // Phase Flip==Z if b=01...0 == 'x', i.e. x[n-2]==1

  // Undo the local registers
  for(j = n-2; j > 0; j--)
    Toffoli(x[j-1], b[j+1], x[j]);
  Toffoli(b[1], b[0], x[0]);

  // Restore b
  for(j = 0; j < n; j++)
    if(j!=1) X(b[j]);
}

/***********************************************
 Squaring a(x) = a0 + a1*x + ... + a_(n-1)*x^(n-1)
 over the ring GF(2)[x] mod (x^n + x + 1).

 Result placed in the n-qubit register b
************************************************/
void Sqr(qbit a[n], qbit b[n]) {

  int i;
  int k;

  // Using forall indicates the CNOT's are independent
  // So these instructions can be done in parallel
  // No forall yet
  // forall(i=0; i<=(n-1)/2; i++) {
  for(i=0; i<=(n-1)/2; i++) {
    k = 2*i;
    CNOT(a[i], b[k]);
  }

  // Would it make a difference to split this into two loops?
  // Maybe since deleaving would be two forall loops!
  // Or perhaps it is okay to just replace the for with a forall.
  for(i=(n+1)/2; i<n; i++) {
    k = (2*i)-n;
    CNOT(a[i], b[k]);
    CNOT(a[i], b[k+1]);
  }
}

/***************************************************
 Program to compute the sqrt(x) in the polynomial ring

     GF(2)[x] / (x^n+x+1)

 Elements of the ring are represented as polynomials
 of maximum degree n-1:

   a(x) = a[0] + a[1]*x + ... + a[n-1]*x^(n-1)

 Use grover search.

****************************************************/
int main() {

  // Quantum registers
  qbit a[n];  // Polynomials to search
  qbit b[n];  // b(x) = a(x)*a(x) mod (x^n + x + 1)
  // qbit t[1];  // Test result if b(x) == x

  // Grover parameters and step index
  int N= pow(2,n);
  int nstep = floor((pi/4)*sqrt(N));
  int istep;

  // Holds final measurement values
  // cbit mt[1]; // measure t
  // cbit ma[n]; // measure a  This holds the square root
  // cbit mb[n]; // measure b  No need to measure since b(x) should be x
  int i;

  // Initialize a[0..n-1] into uniform superposition
  // No forall yet
  // forall(i=0; i<n; i++) H(a[i]);
  for(i=0; i<n; i++) H(a[i]);

  // X(t[0]);
  // H(t[0]);

  qbit x[n-1];  // scratch register
  for(int j=0; j<n-1; j++) PrepZ(x[j],0);

  // Grover iteration: Mark then diffuse
  for(istep=1; istep<=nstep; istep++) {

    Sqr(a, b);         // Sets b(x) = a(x) * a(x)
    EQxMark(b, x);  // Tests if b(x) == x and if so Phase Flips
    Sqr(a, b);         // Note: Sqr is its own inverse

    // Diffuse
    diffuse(a, x);
  }

  // For the final measurement, compute causal state
  Sqr(a, b);
  EQxMark(b, x);   // Note; 1 implies test result b(x)==x is returned in t

  // Now measure and report
  // mt[0] = MeasZ(t[0]);  // If mt[0]==1 then success
  // measure a to recover the square root
  // for(i=0; i<n; i++) ma[i] = MeasZ(a[i]);

  return 0;
}

// Problem instances
//main(2);
//main(3);
//main(5);
//main(8);
//main(13);
//main(19);