// simulation using two quantum bits
//
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <tgmath.h>

typedef double complex Complex;

// magnitude squared
double abs2( Complex c) { double x = creal(c), y = cimag(c); return x*x + y*y; }

// swap a[i] and a[j]
void swap( Complex a[], int i, int j) { Complex t = a[i]; a[i] = a[j]; a[j] = t; }

//    q1 q0
// a0  0 0
// a1  0 1
// a2  1 0
// a3  1 1
//
void X( Complex a[], int k) { // NOT on qubit k
  if( k == 0) { swap( a, 0, 1); swap( a, 2, 3); }
  else        { swap( a, 0, 2); swap( a, 1, 3); }
}
void CX( Complex a[], int c, int k) { // controlled-NOT on qubit k controlled by qubit c
  if( c == 1) swap( a, 2, 3); else swap( a, 1, 3);
}
//
//  0.0       ... r ...                       1.0
//   |<------->|<-------->|<------->|<-------->|
//       p00        p01       p10       p11
//
int Mk( Complex a[]) { // measure overall state
  double r = rand() / (RAND_MAX + 1.0); // 0.0 <= r < 1.0
  double p = 0; // cumulative probability
  int k = 3; for( int i = 0; i < 4; ++i) { p += abs2( a[i]); if( r < p) { k = i; break; } }
  printf( "k = %i\n", k); return k;
}
void M10( Complex a[]) { // measure qubits 1 and 0
  int k = Mk( a); a[0] = a[1] = a[2] = a[3] = 0; a[k] = 1;
}
void M1( Complex a[]) { // measure qubit 1
  int k = Mk( a); if( k == 0 || k == 1) { // qubit 1 is 0
    printf( "M1 = 0\n"); a[2] = a[3] = 0; // normalize remaining states
    double s = sqrt( abs2(a[0]) + abs2(a[1])); a[0] /= s; a[1] /= s; }
  else { // qubit 1 is 1
    printf( "M1 = 1\n"); a[0] = a[1] = 0;
    double s = sqrt( abs2(a[2]) + abs2(a[3])); a[2] /= s; a[3] /= s; }
}
void print( const char msg[], const Complex a[]) {
  printf( "%s:", msg);
  for( int i = 0; i < 4; ++i) printf( " (%g,%g)", creal(a[i]), cimag(a[i]));
  printf( "\n");
}
int main( int argc, char *argv[]) {
  srand(time(0)); double s30 = 1/sqrt(30);
  Complex a[4] = { s30, CMPLX(0,2*s30), 3*s30, CMPLX(0,4*s30) }; print( "a", a);
  X( a, 0); print( "X0", a);
  X( a, 1); print( "X1", a);
  X( a, 1); print( "X1", a);
  X( a, 0); print( "X0", a);
  CX( a, 1, 0); print( "CX10", a);
  CX( a, 0, 1); print( "CX01", a);
  CX( a, 0, 1); print( "CX01", a);
  CX( a, 1, 0); print( "CX10", a);

  // Mk( a);
  // M10( a); print( "M10", a);
  M1( a); print( "M1", a);
}