#include <TSpline.h>
#include <TRandom.h>

#include <iostream>
#include <HEPUnits.hxx>
#include "IDriftManager.hxx"
#include "IFieldManager.hxx"





//***********************************************************
COMET::IDriftManager::IDriftManager() {
//***********************************************************

  fInitialized = false;
     
}

//***********************************************************
COMET::IDriftManager::~IDriftManager() {
//***********************************************************

}


//***********************************************************
void COMET::IDriftManager::InitializeFromDB(bool UseAdaptiveStepSize, double MaxDriftStepTime, double MinDriftStepTime, 
					 int MaxNbrDriftSteps, double ErrorTol, double InitialDriftStepTime, double OmegaTau, double Mu,double DriftVelocity) {
//***********************************************************


  fInitialDriftStepTime = InitialDriftStepTime;
  fUseAdaptiveStepSize = UseAdaptiveStepSize;
  fUseAdaptiveStepSize = true;
  if (fUseAdaptiveStepSize) {
    fMaxDriftStepTime = MaxDriftStepTime;
    fMinDriftStepTime = MinDriftStepTime; 
    //fMinDriftStepTime = MinDriftStepTime; 
    fMaxNbrDriftSteps =  MaxNbrDriftSteps;
   fErrorTol         = ErrorTol;
  } else {
    fMaxDriftStepTime = fInitialDriftStepTime + 1 * unit::ns;
    fMinDriftStepTime = fInitialDriftStepTime - 1 * unit::ns;
    fMaxNbrDriftSteps = 1000;
    fErrorTol = 1000000 * unit::mm;
  }

  //Some constants for the velocity function
  fOmegaTau = OmegaTau;
  fMu = Mu;

  fInitialized = true;
  fCathode = 6.6;
  fDriftVelocity =  DriftVelocity;

}


//***********************************************************
void COMET::IDriftManager::InitializeFromDB(double InitialDriftStepTime, double OmegaTau, double Mu, double DriftVelocity) {
//***********************************************************


  fInitialDriftStepTime = InitialDriftStepTime;
  fUseAdaptiveStepSize = false; 
  fMaxDriftStepTime = fInitialDriftStepTime + 1 * unit::ns;
  fMinDriftStepTime = fInitialDriftStepTime - 1 * unit::ns;
  fMaxNbrDriftSteps = 1000;
  fErrorTol = 1000000 * unit::mm;
  
  //Some constants for the velocity function
  fOmegaTau = OmegaTau;
  fMu = Mu;
  fDriftVelocity =  DriftVelocity;

  fInitialized = true;
  fCathode = 6.6;
}

//***********************************************************
bool COMET::IDriftManager::DriftUntilXLim(IElectronCloud& cloud, double xLim) {
//***********************************************************
  //Dummy variables
  int nbrSteps = 0;
  int nbrStepCalls = 0;
  int nbrVCalls = 0;
  double hMin;

  //Protection to avoid attempts to drift clouds starting in X<0 to RP1 (X>0) or any xLim>0 
  //some classes calling IDriftManager may be doing this kind of attempt as a remain of previous versions 
  if (cloud.GetPosition().X()<0 && xLim>0) {
    std::cout << "WARNING in IDriftManager::DriftUntilXLim: negative value for xLim expected" << std::endl;
    return false;
  }

  //All counter inputs are temporary just to test the adaptive step size. 
  //These counters will be cleaned away in later versions
  return Drift(xLim, true, cloud, nbrSteps, nbrStepCalls, nbrVCalls, hMin);
}

//***********************************************************
bool COMET::IDriftManager::DriftUntilTLim(IElectronCloud& cloud, double tLim) {
//***********************************************************
  //Dummy variables
  int nbrSteps = 0;
  int nbrStepCalls = 0;
  int nbrVCalls = 0;
  double hMin;
  //All counter inputs are temporary just to test the adaptive step size. 
  //These counters will be cleaned away in later versions
  return Drift(tLim, false, cloud, nbrSteps, nbrStepCalls, nbrVCalls, hMin);
}

//***********************************************************
bool COMET::IDriftManager::Drift(double lim, bool driftUntilXLim, IElectronCloud& cloud, 
			      int& nbrSteps, int& nbrStepCalls, int& nbrVCalls, double& hMin) {
//***********************************************************

  if (!fInitialized){
    std::cerr<< "IDriftManager not initialized !!!!" << std::endl;
    return false;
  }

  fNbrStepCalls = 0;
  fNbrVxCalls = 0;
  fMinUsedH = 99999999;
  double meanDriftXVelocity=0;
  nbrSteps = 0;

  bool reachedTheLim = false;
  double h = fInitialDriftStepTime;
  //Save start time
  double startTime = cloud.GetTime();
  double charge = cloud.GetCharge();
  //std::cout << "charge " << charge  << std::endl;
  //Current position and time and velocity
  TLorentzVector xcurr(cloud.GetPosition(), cloud.GetTime());
  TVector3 vcurr(cloud.GetVelocity());

   // Decide if we are moving to larger or smaller values of X (for reverse drift)
  double initsign=0.;
  if (driftUntilXLim) initsign = xcurr.X() - lim; //possibility of going from pad plane to cathode -> Laser
  


//   //TEMP
//   std::cout<<"************************************************************************************"<<std::endl;
//   std::cout<<"************************************************************************************"<<std::endl;
//   std::cout<<"**** New cloud to drift **** X Start at (x,y,z) "<<xcurr.X()<<" "<<xcurr.Y()<<" "<<xcurr.Z()<<std::endl;
//   std::cout<<"************************************************************************************"<<std::endl;
  
//this check for position would return false for reverse drift -> Laser
/*
  //Before we drift check if the current position (or current time) reached the limit already
  if (((driftUntilXLim)&&(fabs(xcurr.X()) > lim)) || ((!driftUntilXLim)&&((xcurr.T() - startTime) > lim))) {
    //Then we return false without setting fDrifted to true
    return false;
  }
*/

 if ((!driftUntilXLim)&&((xcurr.T() - startTime) > lim)) {
    //Then we return false without setting fDrifted to true
    return false;
  }

 //Let the cloud do drift steps until it either exited the drift volume, drifted too long time or it took too many steps
  while ((!reachedTheLim)&&(nbrSteps < fMaxNbrDriftSteps)) {
    //Make a drift step with the current drift step time
    bool useFirstKs = false; //Calculate the first Ks
    TLorentzVector xtemp = Step(xcurr, charge, h, useFirstKs);
  
    if (fUseAdaptiveStepSize) {
  
      //Make two drift steps with half the current drift step time
      useFirstKs = true;
      TLorentzVector xmid = Step(xcurr, charge, h/2, useFirstKs);
      useFirstKs = false; //Calculate the first Ks
      TLorentzVector xhalfh = Step(xmid, charge, h/2, useFirstKs);
      //Get an estimate of the error with step time h
      //double errorEst = (xhalfh.Vect() - xtemp.Vect() ).Mag()/2;
      double errorEst = (xhalfh.Vect() - xtemp.Vect() ).Mag();
     
  
      //If Eest < Etol, or if we use minimal time step accept xtemp
  
      if ((errorEst < fErrorTol)||(h < fMinDriftStepTime)) {
	//std::cout << "error estimate accepted " << std::endl;
	//Before accepting xtemp, save current velocity
	vcurr = (xtemp.Vect() - xcurr.Vect()) * (1/h);
	//Save mean drift x velocity for diffusion
	meanDriftXVelocity += vcurr.X();
	//Now, accept xtemp
	xcurr = xtemp;
	//If this drift step was accepted count it
	nbrSteps++;
	//std::cout << "Number of steps " <<  nbrSteps << " H " << h << std::endl;
	//Save smallest h
	if (h < fMinUsedH) fMinUsedH = h;
	if (h < fMinDriftStepTime) std::cout <<"WARNING in IElectronCloud::Drift: Minimal step size reached" << std::endl;
      }
      //Adjust the step size so that next step get an error about Etol / 2
      //safety factor, 2, prevents infinit loops (the errorEst test is just to not divide by zero)
      if (errorEst > 0.00000001) {
	h = h * pow(fErrorTol / (2 * errorEst), 1./4.);
      } else {
	h = fMaxDriftStepTime;
      }
      //Check so that h didn't pass the limits
      if (h < fMinDriftStepTime) h = fMinDriftStepTime - 0.000001;
      if (h > fMaxDriftStepTime) h = fMaxDriftStepTime;
    } else {
      //Before accepting xtemp, save current velocity
      vcurr = (xtemp.Vect() - xcurr.Vect()) * (1/h);
      //If we don't use adaptive step size it's no question, we will accept xtemp
      xcurr = xtemp;
      //Save mean drift x velocity for diffusion
      meanDriftXVelocity += vcurr.X();
      //Count the steps
      nbrSteps++;
    
    }
    //In case we are drifting positrons it could be possible the cloud reached the central cathode
    //and even went beyond central cathode if the drift time somehow got too big... Then dont use this cloud...
   
    if ((xcurr.X()>0 && cloud.GetStartPosition().X()<=0)||(xcurr.X()<=0 && cloud.GetStartPosition().X()>0)) {
      //TEMP
      std::cout<<"WARNING, cloud reached and passed cathode"<<std::endl;
 
     
      if (cloud.GetStartPosition().X()<=0) fCathode = -fCathode;
     
      double timeCorrection = 0;
  
      if (vcurr.X() != 0) timeCorrection = (fabs(xcurr.X() - fCathode))/fabs(vcurr.X());
 
      //std::cout << " Current X " << xcurr.X() << " current time "  << xcurr.T() << std::endl;
      //std::cout << "fabs " << fabs(xcurr.X() - cathlim) << std::endl;
      //std::cout << " Time correction " <<  timeCorrection << std::endl;
      //std::cout << " fCathode " <<  fCathode << std::endl;
      if (xcurr.X() > 0) xcurr = TLorentzVector(fCathode, xcurr.Y(), xcurr.Z(), xcurr.T() - timeCorrection);
      if (xcurr.X() < 0) xcurr = TLorentzVector(fCathode, xcurr.Y(), xcurr.Z(), xcurr.T() - timeCorrection); 
      
      //std::cout << " Apply correction " << xcurr.X() << " current time "  << xcurr.T() << std::endl;

      reachedTheLim = true;
      //std::cout << " Corrected Current X " << xcurr.X() << " Corrected current time "  << xcurr.T() << std::endl;

      //xcurr.SetX(0.);
      
    
      //If then we don't want to deal with the cloud
      //return false;
    }
    //Check if the cloud went pass the x limit
    double sign = 0;
    sign = xcurr.X()-lim;
   
    if ((driftUntilXLim)&&((sign<0 && initsign>0)||(sign>0 && initsign<0))){ //possibility of reverse drift
      //    if ((driftUntilXLim)&&((fabs(xcurr.X()) > lim))) {
      //If the cloud went further than the readout board the drift is succesfull 
      //and we have to correct if the cloud went too far in the last step
      
      reachedTheLim = true;
      //First correct time
      double timeCorrection = 0;
      if (vcurr.X() != 0) timeCorrection = fabs(xcurr.X() - lim)/fabs(vcurr.X());
      //Then correct the x-position
      //Correct transverse points as well.  
    
      //double yCorr = timeCorrection*vcurr.Y();
      //double zCorr = timeCorrection*vcurr.Z();
          
      //xcurr =  TLorentzVector(lim, xcurr.Y()-yCorr, xcurr.Z()-zCorr, xcurr.T() - timeCorrection);
      xcurr =  TLorentzVector(lim, xcurr.Y(), xcurr.Z(), xcurr.T() - timeCorrection);
      //if (xcurr.X() > 0) xcurr = TLorentzVector(lim, xcurr.Y(), xcurr.Z(), xcurr.T() - timeCorrection);
      //if (xcurr.X() < 0) xcurr = TLorentzVector(-lim, xcurr.Y(), xcurr.Z(), xcurr.T() - timeCorrection);
    
      //Check if the total drift time went past the time limit
    } else if ((!driftUntilXLim)&&((xcurr.T() - startTime) > lim)) {
      //If the limit is time based and the total drift time is more than the time limit, 
      //the drift is succesfull and we have to correct if the cloud used too much time in the last step
      //std::cout << " xcurr " << xcurr.X() << std::endl;
      reachedTheLim = true;
      //First correct the x-position
       
      double dT = (xcurr.T() - startTime) - lim;
      double xCorrection = vcurr.X() * dT;
      //double yCorr = vcurr.Y() * dT;
      //double zCorr = vcurr.Z() * dT;
    
      //Then correct the time
      xcurr = TLorentzVector(xcurr.X() - xCorrection, xcurr.Y(), xcurr.Z(), startTime + lim);
    
    }
  }
  if (nbrSteps >= fMaxNbrDriftSteps) std::cout<<"WARNING! in IDriftManager::Drift: Max number drift steps was reached"<<std::endl;

  nbrStepCalls = fNbrStepCalls;
  nbrVCalls = fNbrVxCalls;
  hMin = fMinUsedH;
  if (nbrSteps > 0) meanDriftXVelocity /= nbrSteps;


  cloud.SetPosition(xcurr.Vect());
  cloud.SetVelocity(vcurr);
  cloud.SetTime(xcurr.T());
  cloud.SetNumberDriftSteps(nbrSteps);
  cloud.SetIsDrifted(true);
  cloud.SetMeanDriftXVelocity(meanDriftXVelocity);



  return reachedTheLim;
}

// TVector3 COMET::IDriftManager::GetRandomArrivalPosition(double tranDiffCoef) {
//   if (IsDrifted()) {
//     //Get the x-displacement due to E-field
//     double totalDx_E = GetXDriftLength();
//     if (totalDx_E > 0.000001) {
//       //Standard deviation 
//       double tranSigma = tranDiffCoef*sqrt( totalDx_E );
//       //Return diffused arrival time around the clouds arrival position (only in Y and Z ofcourse)
//       return TVector3 (fPos.X(), gRandom->Gaus(fPos.Y(), sigma + fYInitialSigma), gRandom->Gaus(fPos.Z(), sigma + fZInitialSigma));
//     }
//   }
//   std::cout<<" in IElectronCloud::GetRandomArrivalTime; Returning (-1,-1,-1)"<<std::endl;
//   return TVector3(-1.,-1.,-1.);
// }


//***********************************************************
TLorentzVector COMET::IDriftManager::Step(const TLorentzVector& x, double charge, double driftTime, bool useFirstKs) {
//***********************************************************

  fNbrStepCalls++;
  TVector3 xstart = x.Vect();
  double   tstart = x.T();
  //Temporary use constant E field and multiply with the -charge (so that electrons (fCharge = -1) goes 
  //in the E direction and positrons (fCharge = 1) goes in the opposit E direction)
  // double Ex = (-charge) * 274.65 * (unit::volt / unit::cm);
  //Use drift velocity and moblility to set E field in simulation. 
 
  double Ex = -fDriftVelocity/fMu;
 

  TVector3 E(0,0,0); //Constant for now... should also be a field map
  //Define a 3D E-field (constant for now). If we are on the -x side of the cathode
  //E field goes in oposit direction
  if (x.X() > 0.0) {
    E = TVector3(Ex,0,0);
  } else if (x.X() < 0.0) {
    E = TVector3( - Ex,0,0);
  } else {
    std::cout<<"IDriftVolume::Step, drift electron started exactly on central cathod and stayed there!"<<std::endl;
    return x;
  }
  
  TVector3 B1(0,0,0);
  if (!useFirstKs) B1 = IFieldManager::GetFieldValue(TLorentzVector(xstart, tstart)); //just to save a call to B
  //std::cout << " B1 " << B1.X() << " " << B1.Y() << " " << B1.Z()  << std::endl;
  //Get the 4 Runge Kutta constants for X
  double k1x = 0;
  if (useFirstKs) {
    k1x = fK1x;
  } else { //Calculate these first Ks for first time and save them
    k1x = vx(charge,E, B1);
    fK1x = k1x;
  }
  //TEMP
  //std::cout<<"vx "<<k1x / (unit::cm / unit::microsecond)<<" us"<<std::endl;

  double x2 = xstart.X()+driftTime/2.*k1x;
  double t2 = tstart+driftTime/2.;
  TVector3 B2x = IFieldManager::GetFieldValue(TLorentzVector(TVector3(x2, xstart.Y(), xstart.Z()), t2));
  double k2x = vx(charge,E, B2x);
  double x3 = xstart.X()+driftTime/2.*k2x;
  double t3 = tstart+driftTime/2.;
  TVector3 B3x = IFieldManager::GetFieldValue(TLorentzVector(TVector3(x3, xstart.Y(), xstart.Z()), t3));
  double k3x = vx(charge,E, B3x);
  double x4 = xstart.X()+driftTime*k3x;
  double t4 = tstart+driftTime;
  TVector3 B4x = IFieldManager::GetFieldValue(TLorentzVector(TVector3(x4, xstart.Y(), xstart.Z()), t4));
  double k4x = vx(charge,E, B4x);
  //Get the new X position
  double xnew = xstart.X() + driftTime/6. * (k1x + 2*k2x + 2*k3x + k4x);
  //Get the 4 Runge Kutta constants for Y
  double k1y = 0;
  if (useFirstKs) {
    k1y = fK1y;
  } else { //Calculate these first Ks for first time and save them
    k1y = vy(charge,E, B1);
    fK1y = k1y;
  }
  double y2 = xstart.Y()+driftTime/2.*k1y;
  TVector3 B2y = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), y2, xstart.Z()), t2));
  double k2y = vy(charge,E, B2y);
  double y3 = xstart.Y()+driftTime/2.*k2y;
  TVector3 B3y = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), y3, xstart.Z()), t3));
  double k3y = vy(charge,E, B3y);
  double y4 = xstart.Y()+driftTime*k3y;
  TVector3 B4y = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), y4, xstart.Z()), t4));
  double k4y = vy(charge,E, B4y);
  //Get the new Y position
  double ynew = xstart.Y() + driftTime/6. * (k1y + 2*k2y + 2*k3y + k4y);
  //Get the 4 Runge Kutta constants for Z
  double k1z = 0;
  if (useFirstKs) {
    k1z = fK1z;
  } else { //Calculate these first Ks for first time and save them
    k1z = vz(charge, E, B1);
    fK1z = k1z;
  }
  double z2 = xstart.Z()+driftTime/2.*k1z;
  TVector3 B2z = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), xstart.Y(), z2), t2));
  double k2z = vz(charge,E, B2z);
  double z3 = xstart.Z()+driftTime/2.*k2z;
  TVector3 B3z = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), xstart.Y(), z3), t3));
  double k3z = vz(charge,E, B3z);
  double z4 = xstart.Z()+driftTime*k3z;
  TVector3 B4z = IFieldManager::GetFieldValue(TLorentzVector(TVector3(xstart.X(), xstart.Y(), z4), t4));
  double k4z = vz(charge,E, B4z);
  //Get the new Z position
  double znew = xstart.Z() + driftTime/6. * (k1z + 2*k2z + 2*k3z + k4z);

  
//   //TEMP
//   std::cout<<"B1 (x,y,z)  "<<B1.X()/unit::tesla<<" "<<B1.Y()/unit::tesla<<" "<<B1.Z()/unit::tesla<<" Tesla"<<std::endl;
//   std::cout<<"v1 (x,y,z)  "<<vx(E,B1)/(unit::mm/unit::second)
//       <<" "<<vy(E,B1)/(unit::mm/unit::second)<<" "<<vz(E,B1)/(unit::mm/unit::second)<<" mm/s"<<std::endl;
//   std::cout<<"deltaX (x,y,z,t) "<<(xnew - x.X())/unit::mm<<" "<<(ynew - x.Y())/unit::mm
//       <<" "<<(znew - x.Z())/unit::mm<<" "<<driftTime<<std::endl;
//   //std::cout<<"Xnew (x,y,z,t) "<<xnew/unit::mm<<" "<<ynew/unit::mm<<" "<<znew/unit::mm<<" "<<tstart+driftTime<<std::endl; 
//   double deltaY = vy(E,B1)*driftTime + (vz(E,B1) * B1.X() - vx(E,B1) * B1.Z()) * 
//     1.60217646 * pow(10.,-19.) * unit::coulomb *
//     driftTime * driftTime / (2 * 9.10938188 * pow(10.,-31.) * unit::kilogram);
// //   double deltaY = vy(E,B1)*driftTime + (vz(E,B1) * B1.X() - vx(E,B1) * B1.Z()) * 
// //     UNO::electron_charge *
// //     driftTime * driftTime / (2 * 9.10938188 * pow(10.,-31.) * unit::kilogram);
//   std::cout<<"delta Y   =  "<<deltaY / unit::mm<<" mm"<<std::endl;
  


  //New position
  return  TLorentzVector(xnew, ynew, znew, tstart + driftTime);
}

// www.pd.infn.it/gruppi/g1/2002Vavra_student_lecture.pdf
//**************************************************
//old drift 
//**************************************************

/*
double COMET::IDriftManager::vx(double charge, TVector3 E, TVector3 B) {

  fNbrVxCalls++;
  if (B.Mag() < 0.00001) return fMu*charge* E.X();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau*charge*charge);
  return con * (E.X() + fOmegaTau*charge/(B.Mag())*(-E.Y()*B.Z() + E.Z()*B.Y()) + fOmegaTau*fOmegaTau*charge*charge*E*B/(B*B)*B.X() );
}

double COMET::IDriftManager::vy(double charge, TVector3 E, TVector3 B) {

  if (B.Mag() < 0.00001) return fMu *charge* E.Y();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau*charge*charge);
 
  double val = con*(E.Y() + fOmegaTau*charge/(B.Mag())*(E.Z()*B.X() + E.X()*B.Z()) -  fOmegaTau*fOmegaTau*charge*charge*E*B/(B*B)*B.Y() );
  return val;
}

double COMET::IDriftManager::vz(double charge, TVector3 E, TVector3 B) {

  if (B.Mag() < 0.00001) return fMu *charge* E.Z();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau*charge*charge);
  double val =  con * (E.Z() + fOmegaTau*charge/(B.Mag())*(-E.X()*B.Y() - E.Y()*B.X()) - fOmegaTau*fOmegaTau*charge*charge*E*B/(B*B)*B.Z() );
  return val;
  
}
*/
  //*********************************************************************
  //Drift Updtate 
  //*********************************************************************
  // should be symmtric in drifting forward and reverse drift.
 
double COMET::IDriftManager::vx(double charge, TVector3 E, TVector3 B) {

  fNbrVxCalls++;
  //double realitycharge = -1.0;
  if (B.Mag() < 0.00001) return fMu*charge* E.X();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau);
  return con * (E.X() - fOmegaTau/(B.Mag())*(E.Y()*B.Z() - E.Z()*B.Y()) + fOmegaTau*fOmegaTau*E*B/(B*B)*B.X() );
}

double COMET::IDriftManager::vy(double charge, TVector3 E, TVector3 B) {

  //double realitycharge = -1.0;
  if (B.Mag() < 0.00001) return fMu *charge* E.Y();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau);
 
  return con*(E.Y() - fOmegaTau/(B.Mag())*(-E.Z()*B.X() - E.X()*B.Z()) +  fOmegaTau*fOmegaTau*E*B/(B*B)*B.Y() );
 
}

double COMET::IDriftManager::vz(double charge, TVector3 E, TVector3 B) {

  if (B.Mag() < 0.00001) return fMu *charge* E.Z();  //Since omega ~ B
  double con = fMu*charge / (1 + fOmegaTau*fOmegaTau*charge*charge);
  return  con * (E.Z() - fOmegaTau/(B.Mag())*(E.X()*B.Y() + E.Y()*B.X()) + fOmegaTau*fOmegaTau*E*B/(B*B)*B.Z() );
 
}