// Print state indices of an n-qubit system,
// sorted by the pairs involved in computations on the k'th qubit.
//
// If compiled using -fopenmp displays parallel processor allocation.
// Specify number of threads using OMP_NUM_THREADS environment variable.
//
// Usage: states [n [k]]
//
// default: n = 6, k = 0
//
// R. Perry, March 2019
//
#include <stdio.h>
#include <stdlib.h>
#ifdef _OPENMP
#include <omp.h>
#endif

// print in binary
//
void printb( unsigned long x, unsigned int nbits)
{
  unsigned long mask = 1LU << (nbits-1);

  while( mask) { printf( " %c", (mask & x) ? '1' : '0'); mask >>= 1; }

  putchar( '\n');
}

// linear start,stop thread assignment
//
void linear( unsigned long ID, unsigned long start, unsigned long stop, unsigned long k1)
{
  unsigned long count = 0, i = start, j;

#ifdef BAD
  i += (start & k1); // start index could be bad, adjust it
#endif

  printf( "proc: %lu start: %lu i: %lu stop: %lu\n", ID, start, i, stop);

  while( i < stop)
  {
    j = i | k1;  printf( "proc: %lu i: %lu j: %lu\n", ID, i, j);

    ++count;  ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }

  printf( "proc: %lu count: %lu\n", ID, count);
}

int main( int argc, char *argv[])
{
  unsigned int n = 6;  if( argc > 1) n = atoi(argv[1]);

  unsigned int k = 0;  if( argc > 2) k = atoi(argv[2]);

  if( k >= n) { fprintf( stderr, "%s: k < n required\n", argv[0]); return 1; }

  printf( "%u qubits, k = %u\n\n", n, k);

 // k1 = 0...010...0, k'th bit = 1
 //
 // i will have k'th bit = 0, j will have k'th bit = 1
 //
  unsigned long k1 = 1LU << k;

#ifdef _OPENMP

  unsigned long T = omp_get_max_threads(); printf( "T: %lu\n", T); // OMP_NUM_THREADS

#ifdef BAD
//
// linear start,stop thread assignment
// half of the processors won't be utilized when k >= b
//
  unsigned long N = 1LU << n;

 // check for overflow
 //
  unsigned long long check = (unsigned long long)T*N; check /= T;

  if( check != N) { fprintf( stderr, "overflow\n"); exit(1); }

  #pragma omp parallel num_threads(T)
  {
    unsigned long ID = omp_get_thread_num();
    unsigned long start = ((unsigned long long)ID*N)/T;
    unsigned long stop = (((unsigned long long)ID+1)*N)/T;

    linear( ID, start, stop, k1);
  }

#else // not BAD

 // T = 2**t 
 //
  unsigned int t = 0;  while( (1LU << t) < T) ++t;

  if( (1LU << t) != T) { fprintf(stderr,"number of threads is not a power of 2\n"); exit(1); }

  if( n <= t) { fprintf( stderr, "too many threads, not enough qubits\n"); exit(1); }

  unsigned int b = n-t; // low-order bits of amplitude index

  if( k < b)
  #pragma omp parallel num_threads(T)
  {
   // ID*N/T = ID*(2**n)/(2**t) = ID*2**(n-t)
   //
    unsigned long ID = omp_get_thread_num(), start = ID << b, stop = (ID+1) << b;

    linear( ID, start, stop, k1);
  }
  else // k >= b
  {
    unsigned int d = k-b; // low-order bits of ID, d >= 0

    #pragma omp parallel num_threads(T)
    {
      unsigned long ID = omp_get_thread_num();

     // bit layout:
     //     <------------------------------- n ----------------------------------->
     //     <------ t-d-1 ----------><- 1 -><----- d ------><--- 1 ----><-- b-1 -->
     // i = top (t-d-1) bits from ID, bit k, d bits from ID, bit 0 of ID,
     //     ^^^^^^^^^^^^^^^^^^^^^^^^         ^^^^^^^^^^^^^^
     // one of the indicated pieces could be empty depending on k
     //
      unsigned long i =   ((ID >> (d+1)) << (k+1))
			| (((ID >> 1) & ((1LU << d) - 1)) << b)
			| ((ID & 1) << (b-1));

      unsigned long count = 0, L = 1LU << (b-1), j;

      while( count < L)
      {
        j = i | k1;  printf( "proc: %lu i: %lu j: %lu\n", ID, i, j);

        ++count;  ++i; // count will be L after this loop
      }

      printf( "proc: %lu count: %lu\n", ID, count);
    }
  }

#endif // BAD

#else // not _OPENMP

  unsigned long N = 1LU << n, i = 0, j;

  while( i < N)
  {
    j = i | k1;  printb( i, n); printb( j, n); putchar('\n');

    ++i;  i += (i & k1); // skip i values with k'th bit = 1
  }

#endif // _OPENMP

  return 0;
}