// quantum computing emulation: CORVIDS
//
// Grover search to find c such that all elements of y+c*R are non-negative
//
// References:
//  Complete recovery of values in Diophantine systems (CORVIDS), Sean Wilner,
//   Katherine Wood, and Daniel Simons, July 02, 2018, https://psyarxiv.com/7shr8
//  The R matrix here is the bottom part of reduced R (HRbot) as described in
//  Quantum CORVIDS, R. Perry, July 2018, http://fog.misty.com/perry/qc/corvids/notes.html
//
// Usage: corvids [-drsSE] [k] < file
//
//   k = number of iterations
//  -d  extra debug output
//  -r  use random global phase factor
//  -s  sort states by probability when printing
//  -S  track sum of solutions
//  -E  show entanglement
//
//  if k is not specified it is set to round((pi/4)*sqrt(q.N/M))
//
// file: n d p M y R
//
//   n = total number of responses
//   d = number of unknown coefficients (c values)
//   p = length of row vector y
//   M = number of solutions
//   y = bottom part of one solution
//   R = d-by-p matrix
//
// R. Perry, July 2018
//
#include <stdio.h>
#include <stdlib.h> // atoi(), rand(), free()
#include <tgmath.h>
#include "qce.h"

unsigned long mask; // b 1's

// struct for sorting by probability
//
typedef struct
{
  unsigned long i;  double v; // index and value
}
P;

// comparison function for qsort
//
int compare( const void *v1, const void *v2)
{
  const P *p1 = v1, *p2 = v2;

  if( p1->v < p2->v) return -1;
  else if( p1->v > p2->v) return +1;
  else return 0;
}

// print state
//
// q represents the concatenation of d b-bit coefficient values
//
void print_q( const char *msg, State q, unsigned int d, unsigned int b, int sort)
{
  static P *p = NULL;  static unsigned long pn = 0;

  printf( "%s%s:\n", msg, sort ? " (sorted)" : "");

  if( sort)
  {
    if( pn != q.N) { pn = q.N; free(p); p = emalloc(pn*sizeof(P)); }

    for( unsigned long i = 0; i < q.N; ++i) { p[i].i = i; p[i].v = abs2(q.a[i]); }

    qsort( p, pn, sizeof(P), compare);
  }

  for( unsigned long i = 0; i < q.N; ++i)
  {
    unsigned long t, k;

    if( sort) t = p[i].i; else t = i; // extract and print the c values

    for( unsigned int j = 0; j < d; ++j) { printf( " %3lu", t & mask); t >>= b; }

    if( sort) k = p[i].i; else k = i;

    printf( "   %g (%g,%g)\n", abs2(q.a[k]), creal(q.a[k]), cimag(q.a[k]) );
  }
}

// check if state t is a valid target
//
// (could calculate which states are valid in advance just once)
//
int valid( unsigned long t, unsigned int d, unsigned int b, unsigned int p, int y[p], int R[d][p])
{
  int x[p]; for( unsigned int k = 0; k < p; ++k) x[k] = y[k];

  for( unsigned int j = 0; j < d; ++j) // x = y + c*R
  { 
    int c = t & mask; for( unsigned int k = 0; k < p; ++k) { x[k] += c*R[j][k]; } t >>= b;
  }

  int nonneg = 1; for( unsigned int k = 0; k < p; ++k) if( x[k] < 0) { nonneg = 0; break; }

  return nonneg;
}

// flip sign of target states
//
void f( State q, unsigned int d, unsigned int b, unsigned int p, int y[p], int R[d][p])
{
  for( unsigned long i = 0; i < q.N; ++i) if( valid( i, d, b, p, y, R)) q.a[i] = -q.a[i];
}

// print sum and average
//
void print_stats( unsigned long M, unsigned int d, unsigned long qsum[d], double psum)
{
  printf( "sum:\n");  for( unsigned int j = 0; j < d; ++j) printf( " %3lu", qsum[j]);

  printf( "   %g\navg:\n", psum);  for( unsigned int j = 0; j < d; ++j) printf( " %g", (double)qsum[j]/M);

  putchar('\n');
}

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

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

  unsigned int seed = SRAND();

 // debug, random phase, and sum options
 //
  if( argc > 1 && argv[1][0] == '-') 
  {
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i])
    {
      case 'd': ++debug; break;
      case 's': ++sort; break;
      case 'S': ++sum; break;
      case 'r': ++rphase; break;
      case 'E': ++Eflag; break;
      default: error( "corvids: bad option");
    }
    --argc, ++argv;
  }

 // input data
 //
  unsigned int n, d, p;  unsigned long M;

  if( scanf("%u%u%u%lu", &n, &d, &p, &M) != 4) error( "corvids: bad input n,d,p,M");

  int y[p], R[d][p];  unsigned long qsum[d];  double psum = 0;

  for( unsigned int i = 0; i < p; ++i) if( scanf("%i",y+i) != 1) error( "corvids: bad input y");

  for( unsigned int i = 0; i < d; ++i) for( unsigned int j = 0; j < p; ++j)
    if( scanf("%i",R[i]+j) != 1) error( "corvids: bad input R");

 // number of qubits per coefficient is smallest b such that 2**b > n
 //
  unsigned int b = 1;  unsigned long B = 1LU << b;  while( B <= n) { ++b; B <<= 1; }

  State q = Q(d*b);  mask = (1LU << b) - 1; // b 1's

// note that the auxiliary qubit is not included here, it is not needed
// since we will just directly flip the sign of states in the iterations

 // number of iterations
 //
  int k; if( argc > 1) k = atoi(argv[1]); else k = (pi/4)*sqrt((double)q.N/M) + 0.5;

  printf( "corvids: seed = %u, b = %u, q.n = %u, q.N = %lu, M = %lu,"
    " n = %u, d = %u, p = %u, k = %i, sort = %d, sum = %d\n",
    seed, b, q.n, q.N, M, n, d, p, k, sort, sum);

  printf( "corvids: y ="); for( unsigned int i = 0; i < p; ++i) { printf( " %3d", y[i]); } putchar('\n');

  printf( "corvids: R ="); for( unsigned int i = 0; i < d; ++i)
    { for( unsigned int j = 0; j < p; ++j) { printf( " %3d", R[i][j]); } if( i < d-1) putchar(','); }
    putchar( '\n');

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

  if( debug) { print_q( "init", q, d, b, sort);  if( Eflag) print( "E", q, PRINT_TANGLE); }

  for( unsigned int i = 0; i < q.n; ++i) H( q, i);

  if( debug) { print_q( "H", q, d, b, sort);  if( Eflag) print( "E", q, PRINT_TANGLE); }

 // iterations
 //
  for( int iter = 1; iter <= k; ++iter)
  {
    if( debug) printf( "iter=%i\n", iter);

   // apply f, i.e. flip sign of target states
   //
    f( q, d, b, p, y, R);

    if( debug) { print_q( "f", q, d, b, sort);  if( Eflag) print( "E", q, PRINT_TANGLE); }

   // inversion about the average
   //
    double complex 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];

    // printf( "avg = (%g,%g)\n", creal(avg), cimag(avg));

    if( debug) { print_q( "inv", q, d, b, sort);  if( Eflag) print( "E", q, PRINT_TANGLE); }

    if( sum) // set sum values and measurement probability
    {
      psum = 0;  for( unsigned int j = 0; j < d; ++j) qsum[j] = 0;

      for( unsigned long i = 0; i < q.N; ++i) if( valid( i, d, b, p, y, R))
      {
        unsigned long t = i;  for( unsigned int j = 0; j < d; ++j) { qsum[j] += t & mask;  t >>= b; }

        psum += abs2(q.a[i]);
      }

      if( debug) print_stats( M, d, qsum, psum);

    } // sum

  } // iterations

  if( !debug) print_q( "final", q, d, b, sort);

  if( !debug && sum) print_stats( M, d, qsum, psum);

  return 0;
}