// qfunct.h  basic functions bra ket unitary algorithms
//           new: many includes and types defined here
#include <complex> // C++ templates 
#include <vector>
#include <iostream>
#include <math.h>
using namespace std;

typedef complex<double> cmplx;
typedef vector<cmplx> c_vec;
typedef vector<c_vec>  c_mat; // defined as vector of complex vectors

static c_vec operator+(c_vec X, c_vec Y)
{
  int n = X.size();
  c_vec Z(n);
  for(int j=0; j<n; j++)
    Z[j] = X[j]+Y[j];
  return Z;
} // operator+

static c_vec operator*(double a, c_vec X)
{
  int n = X.size();
  c_vec Z(n);
  for(int j=0; j<n; j++)
    Z[j] = a*X[j];
  return Z;
}

static void prtmat(c_mat mat)
{
  for(int i=0; i<mat.size(); i++)
    for(int j=0; j<mat[i].size(); j++)
      cout << "mat[" << i << "][" << j << "] = " << mat[i][j] << endl;

} // end prtmat

static void prtvec(c_vec vec)
{
  for(int j=0; j<vec.size(); j++)
     cout << "vec[" << j << "] = " << vec[j] << endl;
} // end prtvec

static cmplx dotp(c_vec X, c_vec Y) // c = X*Y with conj
{
  int n = X.size();
  cmplx Z = cmplx(0.0,0.0);

  for(int i=0; i<n; i++)
    Z = Z + X[i]*conj(Y[i]);
  return Z;
} // end dotp

static c_vec outp(c_vec X, c_vec Y) // Z = X*Y with conj 
{
  int nx = X.size();
  int ny = Y.size();
  c_vec Z(nx*ny);

  for(int i=0; i<nx; i++)
    for(int j=0; j<ny; j++)
      Z[i*ny+j] = X[i]*conj(Y[j]);
  return Z;
} // end outp

static c_vec matvec(c_mat A, c_vec X)
{
  int n = A.size();
  c_vec Y(n);

  for(int i=0; i<n; i++)
  {
    Y[i] = cmplx(0.0,0.0);
    for(int j=0; j<n; j++)
      Y[i] = Y[i] + A[i][j]*X[j];
  }
  return Y;
} // end matvec

static c_vec make_reg(c_vec keta, c_vec ketb)
{
  c_vec Y(4);
  Y[0] = keta[0];
  Y[1] = keta[1];
  Y[2] = ketb[0];
  Y[3] = ketb[1];
  return Y;
}

static c_vec vecmat(c_vec X, c_mat A)
{
  int n = A.size();
  c_vec Y(n);

  for(int i=0; i<n; i++)
  {
    Y[i] = cmplx(0.0,0.0);
    for(int j=0; j<n; j++)
      Y[i] = Y[i] + X[i]*A[i][j];
  }
  return Y;
} // end vecmat

static c_mat tensor(c_vec X, c_vec Y) // A = X*Y with conj 
{
  int nx = X.size();
  int ny = Y.size();
  c_mat A;
  c_vec V;

  for(int j=0; j<ny; j++)
    V.push_back(cmplx(0,0));
  for(int i=0; i<ny; i++)
    A.push_back(V);

  for(int i=0; i<nx; i++)
    for(int j=0; j<ny; j++)
      A[i][j] = X[i]*conj(Y[j]);
  return A;
} // end tensor

static c_vec get_ketp(int n, c_vec kin) // extract ket n 0..(m-2)/2 from kin
       // only works for pure ket, not superposition
{
  int m = kin.size();
  c_vec Y(2); // returned ket
  int i, j, k;
  c_vec ket0(2);
  c_vec ket1(2);

  Y[0] = cmplx(0); // error return
  Y[1] = cmplx(0);
  ket0[0] = cmplx(1);
  ket0[1] = cmplx(0);
  ket1[0] = cmplx(0);
  ket1[1] = cmplx(1);
   
  j = -1;
  for(i=0; i<m; i++) 
  {
    if(kin[i]==1.0)
    {
      j = i;
      break;
    }
  }
  if(j==-1)
  {
    cout << "get_ket error, not pure or bad" << endl;
    return Y;
  }
  for(i=0; i<n; i++) // shift to find bit
  {
    j = j/2;
  }
  if((j&1) == 0) // test bottom bit after shifting
  {
    Y = ket0;
  }
  else if((j&1) == 1)
  {
    Y = ket1;
  }
  return Y;
} // end get_ketp

static c_vec get_ket(int n, c_vec kin) // extract ket n, first input control
{
  int m = kin.size();
  c_vec Y(2); // returned ket

  if(n>m/2) cout << "get_ket n=" << n << " too big" << endl;
  Y[0] = kin[n*2];
  Y[1] = kin[n*2+1];
  return Y;
} // end get_ket

static c_vec cRk(c_vec ctrl, c_vec kin, c_mat Rk) // controlled rotation
{
  c_vec Y(2); // returned ket

  Y = kin;
  if(ctrl[0]!=cmplx(1,0) && ctrl[1]!=cmplx(0,0)) Y = matvec(Rk, kin);
  return Y;
} // end cRk2

// end qfunct.cpp


static void prtket1(c_vec vec) // print a single ket (no << endl;)
{
  double sml = 1.0e-12; // close to zero for printout
  double ar = real(vec[0]);
  double ai = imag(vec[0]);
  double br = real(vec[1]);
  double bi = imag(vec[1]);
  if(abs(ai)<sml & abs(vec[1])<sml)
  {
    cout << "|0>";
  }
  else if(abs(vec[1])<sml)
  {
    cout << "(" << ar << "," << ai << ") |0>";
  }
  else if(abs(bi)<sml & abs(vec[0])<sml)
  {
    cout << "|1>";
  }
  else if(abs(vec[0])<sml)
  {
    cout << "(" << br << "," << bi << ") |1>";
  }
  else if(abs(ai)<sml & abs(bi)<sml)
  {
    cout << "(" << ar << " |0>" << " + "
         << "(" << br << " |1>";
  }
  else 
  {
    cout << "(" << ar << "," << ai << ") |0>" << " + "
         << "(" << br << "," << bi << ") |1>";
  }
} // end prtket1

static void prtket(c_vec vec) // print a single ket (with << endl;)
{
  prtket1(vec);
  cout << endl;
} // end prtket

static void prtket2(c_vec a, c_vec b) // two kets (no << endl;)
{
  prtket1(a);
  prtket1(b);
} // end prtket

static void prtket3(c_vec a, c_vec b, c_vec c) // two kets (no << endl;)
{
  prtket1(a);
  prtket1(b);
  prtket1(c);
} // end prtket


// factor state Y knowing |a>
static c_vec  factor(c_vec Y, c_vec a) // Y = |a>|b>, return b 
{
  double sml = 1.0e-12; // close to zero        
  c_vec b(2);           // returned ket

  if(abs(a[0]) > sml)
  {
    b[0] = conj(Y[0] / a[0]);
    b[1] = conj(Y[1] / a[0]);
  }
  else if(abs(a[1]) > sml)
  {
    b[0] = conj(Y[2] / a[1]);
    b[1] = conj(Y[3] / a[1]);
  }
  else
  {
    cout << "factor fail, bad first argument" << endl;
    prtvec(a);
    cout << endl;
    return a;
  }
  return b;
} // end factor

// factor state Y = |a>|b>|c> knowing |a> and |b>
static c_vec  factor3(c_vec Y, c_vec a, c_vec b) // Y = |a>|b>|c>, return c 
{
  double sml = 1.0e-12; // close to zero        
  c_vec c(2);           // returned ket
  c_vec Y2(4);          // after a reduction

  if(abs(a[0]) > sml)
  {
    Y2[0] = conj(Y[0] / a[0]);
    Y2[1] = conj(Y[1] / a[0]);
    Y2[2] = conj(Y[2] / a[0]);
    Y2[3] = conj(Y[3] / a[0]);
  }
  else if(abs(a[1]) > sml)
  {
    Y2[0] = conj(Y[4] / a[1]);
    Y2[1] = conj(Y[5] / a[1]);
    Y2[2] = conj(Y[6] / a[1]);
    Y2[3] = conj(Y[7] / a[1]);
  }
  else
  {
    cout << "factor3 fail, bad a argument" << endl;
    prtvec(a);
    cout << endl;
    return a;
  }

  if(abs(b[0]) > sml)
  {
    c[0] = Y2[0] / b[0];
    c[1] = Y2[1] / b[0];
  }
  else if(abs(b[1]) > sml)
  {
    c[0] = Y2[2] / b[1];
    c[1] = Y2[3] / b[1];
  }
  else
  {
    cout << "factor3 fail, bad b argument" << endl;
    prtvec(b);
    cout << endl;
    return b;
  }
  return c;
} // end factor3