// quantum computing emulation: wordle, using Grover's search
//
// Usage: wordle [-diMOsS1] [k [m]]
//
//  -d  extra debug output, useful with -s option
//  -i  create image files in TMP/ subdirectory
//  -M  minimize number of states
//  -O  optimize number of iterations, only implemented for the small example
//  -s  small example
//  -S  Sequential, find each letter separately
//  -1  amplify words with at least 1 matching letter (not helpful)
//  -11 -1 just once 
//   k = number of iterations, default = (pi/4)*sqrt(N)
//   m = value to match (sorted word answers index)
//
//  if m is negative or not specified then it is generated randomly
//
// R. Perry, May 2022
//
#include <iostream>
#include <cmath>
#include <cstdio>
#include "qce.h"
#include "words.h" // Nletters, Nans_{small,wordle}, Nwords_{small,wordle}, words_{small,wordle}[]
#include "pgm.h" // {NUMS_}WIDTH1, {NUMS_}HEIGHT, {NUMS_}LENGTH, a2z[][], nums[][]

unsigned long Nans = Nans_wordle, Nwords = Nwords_wordle; const char **words = words_wordle;

using namespace qce;
using namespace std;

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

#define WIDTH (WIDTH1*Nletters)
Real image[HEIGHT][WIDTH]; // grayscale words 

// zero out a section of the image
//
void image_zero( int offset, int width)
{
  for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < width; ++j) image[i][j+offset] = 0;
}

// scale a section of the image to max 255
//
void image_scale( int offset, int width)
{
  Real max = 0; // min will be 0
  for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < width; ++j) {
    Real v = image[i][j+offset]; if( v > max) max = v;
  }
 // cout << "image_scale: max = " << max << '\n';
  Real scale = 255/max;
  for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < width; ++j) image[i][j+offset] *= scale;
}

// image is superposition of all words weighted by their probabilities
//
void image_create( const State &q)
{
  image_zero( 0, WIDTH);

  for( unsigned long k = 0; k < Nwords; ++k) {
    Real p = abs2(q.a[k]); int offset = 0; const char *w = words[k];
    while( *w) { const unsigned char *a = a2z[(*w-'a')*HEIGHT]; // assume ASCII
      for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < WIDTH1; ++j)
        image[i][j+offset] += p*(255-*a++); // invert intensity
      offset += WIDTH1; ++w; // next letter
    }
  }
  image_scale( 0, WIDTH);
}

// update one letter position in the image
//
void image_update( const State &q, int letter)
{
  int offset = letter*WIDTH1; image_zero( offset, WIDTH1);

  for( int k = 0; k < 26; ++k) {
    Real p = abs2(q.a[k]); const unsigned char *a = a2z[k*HEIGHT];
    for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < WIDTH1; ++j)
      image[i][j+offset] += p*(255-*a++); // invert intensity
  }
 image_scale( offset, WIDTH1);
}

// with -S option, image is initially superposition of all letters in each position
//
void image_init( const State &q)
{
  for( int letter = 0; letter < Nletters; ++letter) image_update( q, letter);
}

// insert iteration number and write image to PGM file
//
void image_print( int pgm, int iter, int letter)
{
  char fname[20]; sprintf( fname, "TMP/%03d.pgm", pgm);
  char number[20]; sprintf( number, "%03d", iter);
  FILE *fp = fopen( fname, "w"); if( !fp) { perror(fname); return; }
  fprintf( fp, "P5 %i %i 255\n", WIDTH, HEIGHT);
  int offset = letter*WIDTH1; char *w = number; // insert iteration number at top left of letter
  while( *w) { const unsigned char *a = nums[(*w-'0')*NUMS_HEIGHT];
    for( int i = 0; i < NUMS_HEIGHT; ++i) for( int j = 0; j < NUMS_WIDTH1; ++j)
      image[i][j+offset] += 255-*a++; // invert intensity
    offset += NUMS_WIDTH1; ++w; // next digit
  }
  for( int i = 0; i < HEIGHT; ++i) for( int j = 0; j < WIDTH; ++j) {
    unsigned char v = 255-image[i][j]; // reinvert intensity
    putc( v, fp); // fprintf( fp, "%u\n", v);
  }
  fclose(fp);
}

// flip states with at least 1 matching letter
// this is not helpful, screws up the results
//
void f1( const State &q, unsigned long m)
{
  unsigned long count = 0;
  for( unsigned long k = 0; k < Nwords; ++k) {
    for( const char *ans = words[m], *w = words[k]; *ans; ++ans, ++w)
      if( *ans == *w) { ++count; q.a[k] = -q.a[k]; break; }
  }
  cout << "f1: count = " << count << '\n';
}

// inversion about the average
//
void inv( State &q)
{
  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];
}

// display solution and top two letter or word probabilities
//
void check( const State &q, unsigned long m, int iter, int Sflag)
{
  unsigned long s[2]; Real p[2], t; s[0]=s[1]=0; p[0]=p[1]=abs2(q.a[0]);
  for( unsigned long i = 1; i < q.N; ++i) { t = abs2(q.a[i]);
    if( t > p[0]) { p[1]=p[0]; p[0]=t; s[1]=s[0]; s[0]=i; }
    else if( t > p[1]) { p[1]=t; s[1]=i; }
  }
  Real pm = abs2(q.a[m]); cout << iter << " " << pm << " " << 1-pm << " " << m;
  for( int i = 0; i < 2; ++i) {
    cout << " " << p[i] << " " << 1-p[i] << " " << s[i] << " ";
    if( Sflag) cout << ((s[i] < 26) ? (char)('a'+s[i]) : '*');
     else cout << ((s[i] < Nwords) ? words[s[i]] : "*****");
  }
  cout << '\n';
}

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

  int debug = 0, iflag = 0, Mflag = 0, Oflag = 0, sflag = 0, Sflag = 0, f1flag = 0;

  if( argc > 1 && argv[1][0] == '-') { // options
    for( int i = 1; argv[1][i]; ++i) switch( argv[1][i]) {
      case 'd': ++debug; break;
      case 'i': ++iflag; break;
      case 'M': ++Mflag; break;
      case 'O': ++Oflag; break;
      case 's': ++sflag; break;
      case 'S': ++Sflag; break;
      case '1': ++f1flag; break;
      default: error( "wordle: bad option");
    }
    --argc, ++argv;
  }

  if( Mflag && Oflag) error( "wordle: -M and -O not compatible");

  if( sflag) { Nans = Nans_small; Nwords = Nwords_small; words = words_small; }

 // number of qubits is smallest q.n such that N = 2**q.n >= number of letters or words
 //
  unsigned int n = 1; unsigned long N = 1LU << n, M = Sflag ? 26 : Nwords;

  while( N < M) { ++n; N <<= 1; }

  State q(n); if( Mflag) q.N = M; // if minimized only first M states will be non-zero

  Real a = 1/sqrt(q.N); // for equal amplitudes

 // number of iterations
 //
  int k = (q.pi/4)*sqrt(q.N); if( argc > 1) k = atoi(argv[1]);
 
 // value to match
 //
  unsigned long m; long arg = -1; if( argc > 2) arg = atol(argv[2]);

  if( arg < 0) m = irand(Nans); else m = arg;

  if( m >= Nans) error( "wordle: bad m arg");

  cout << "wordle: Nletters = " << Nletters << ", Nans = " << Nans << ", Nwords = " <<
    Nwords << ", seed = " << seed << ", q.n = " << q.n << ", q.N = " << q.N << 
    ", k = " << k << ", m = " << m << ", word = " << words[m] << "\n";

  if( Sflag) // Sequential, find each letter separately
  {
    for( int pgm = 0, letter = 0; letter < Nletters; ++letter)
    {
     // initial state |1...1>/sqrt(N) = H|0...0>, or |1...10...0>/sqrt(M) if minimized
     //
      q.init(0); for( unsigned long i = 0; i < q.N; ++i) q.a[i] = a;

      if( debug) q.print( "init");

      cout << "letter " << letter << '\n'; int ans = words[m][letter] - 'a';

      check( q, ans, 0, Sflag); if( iflag && letter == 0) { image_init(q); image_print(pgm++,0,letter); }

      for( int iter = 1; iter <= k; ++iter, ++pgm)
      {
        q.a[ans] = -q.a[ans]; if( debug) q.print( "f"); // apply f, i.e. flip sign of ans
    
        inv(q); if( debug) q.print( "inv"); // inversion about the average
    
        check( q, ans, iter, Sflag); if( iflag) { image_update(q,letter); image_print(pgm,iter,letter); }
      }
    }
  }
  else if( sflag && Oflag) // small example with optimal number of iterations (2)
  {
   // initial state |aaaaaaab> with a = sin(t/2), where t satisfies (k+0.5)*t = pi/2,
   //  Nwords = 7, q.N = 8
   //
    a = sin( ((q.pi/2)/(2+0.5))/2 ); Real b = sqrt(1-Nwords*a*a);

    q.init(0); q.a[Nwords] = b; for( unsigned long i = 0; i < Nwords; ++i) q.a[i] = a;

    State s(q.n); memcpy( s.a, q.a, q.N*sizeof(Complex)); // copy of initial state

    if( debug) { q.print( "init"); s.print("s"); }

    check( q, m, 0, Sflag);
  
    for( int iter = 1; iter <= k; ++iter)
    {
      q.a[m] = -q.a[m]; // apply f, i.e. flip sign of state m
  
      if( debug) q.print( "f");
  
     // apply general operator 2|s><s|-I instead of inversion about the average
     //
      Real prod = 0; for( unsigned long i = 0; i < q.N; ++i) prod += s.a[i]*q.a[i];

      prod *= 2; for( unsigned long i = 0; i < q.N; ++i) q.a[i] = s.a[i]*prod-q.a[i];

      if( debug) q.print( "inv");
  
      check( q, m, iter, Sflag);
    }
  }
  else // Parallel, search whole words
  {
   // initial state |1...1>/sqrt(N) = H|0...0>, or |1...10...0>/sqrt(M) if minimized
   //
    q.init(0); for( unsigned long i = 0; i < q.N; ++i) q.a[i] = a;

    if( debug) q.print( "init");

    check( q, m, 0, Sflag); if( iflag) { image_create(q); image_print(0,0,0); }
  
    for( int iter = 1; iter <= k; ++iter)
    {
      if( f1flag) { f1(q,m); if( f1flag > 1) f1flag = 0;  }
       else q.a[m] = -q.a[m]; // apply f, i.e. flip sign of state m
  
      if( debug) q.print( "f");
  
      inv(q); if( debug) q.print( "inv"); // inversion about the average
  
      check( q, m, iter, Sflag); if( iflag) { image_create(q); image_print(iter,iter,0); }
    }
  }

  return 0;
}