#include "qobjects.h"

#include <cstdlib>
#include <ctime>

//------------------------------------------------------------
// This code uses the qobjects set of classes to integrate a 
// a stochastic Master equation describing the continuous 
// measurement of the computational basis of a qubit (or, 
// equivalently, the spin-z component of a spin-1/2 particle).
// The qubit is initialy completely mixed. The code writes
// out the impurity (1-Tr(rho^2)) as the measurement proceeds, 
// averaged over some number of trajectories if desired.
//------------------------------------------------------------

int main(int argc, char* argv[]){
   
//---------------------------------------initialize parameters

   const int seed = 123456;              // random seed
   const std::complex<double> ei(0,1);   // imaginary unit 
   const unsigned int N = 2,             // size of quantum system
                      ns = 1000,         // number of time steps
                      nt = 10;           // number of trajectories
   const double T = 1,                   // total time
                dt = T/ns,
                rdt = std::sqrt(dt), 
                k = 1;                   // measurement rate

//--------------------------------initialise files

   std::ofstream avseries("avImp.dat");
   
//--------------------------------initialise states

   qmatrix fullmix("maxmix",N); 
      
//--------------------------------initialize operators
                    
   qoperator Q("lineig",N);
   qoperator Q2 = Q*Q;
	     
//---------------------------------generate Wiener noise (dW's)

   std::vector<double> noise(ns*nt);
   make_noise(noise); 
   
//------------------------------------------------------------
//                        Simulation
//------------------------------------------------------------

   double t = 0,                      // initialise variables
          dW = 0;
          
   std::complex<double> d_rec = 0;
   std::vector<double> avImp(ns+1);
	 
   time_t time0, *time_ptr = 0;       // timing initialisation
   time0 = time(time_ptr);

   
   for (unsigned int j=0;j<nt;j++) {               // average multiple trajectories
   
     qmatrix rho = fullmix;
     avImp[0] += std::real(1.0 - trace(rho*rho));
   
     for (unsigned int i=0;i<ns;i++) {             // integrate a single trajectory
   
       dW = rdt*noise[i+j*ns];

       d_rec = 4*k*dt*trace(Q*rho) + std::sqrt(2*k)*dW;
       rho  += (-k*dt)*(Q2*rho + rho*Q2 + (-2.0)*Q*rho*Q) 
             + d_rec*(Q*rho + rho*Q)
             + (k*dW*dW - k*dt)*(Q2*rho + rho*Q2 + 2.0*Q*rho*Q);     // Milstien

       rho.normalise();
                 
       avImp[i+1] += std::real(1.0 - trace(rho*rho));
     
     }

     std::cout << "finished traj no. " << j+1 << std::endl;
     
   }
   
   for (unsigned int i=0;i<ns+1;i++) {avImp[i] /= nt;}
   
   
   //-----------------------
   // output simulation time
   //-----------------------

   std::cout << std::endl << "Simulation took " << difftime(time(time_ptr),time0) 
                          << " seconds" << std::endl;
   	              
   //---------------
   // output results
   //---------------
   
   for (unsigned int i=0;i<ns+1;i++) { 
     avseries << i*dt << " " << avImp[i] << std::endl; 
   } 
   avseries.close();
   
   return 0;
}