#include "ClassHelio.h"
#include "TROOT.h"
#include "TBenchmark.h"
#include "TLorentzVector.h"
#include "TMath.h"
#include "TFile.h"
#include "TF1.h"
#include "TTree.h"
#include "TRandom.h"
#include "TClonesArray.h"
#include <vector>
#include <fstream>

//----------- usage 
// $root sim.C+ | tee output.txt
// this will same the massage to output.txt

void knockout(){

   //================================================= User Setting
   //---- reaction
   int AA = 23, ZA = 9;
   int Aa = 1,  Za = 1;
   int A2 = 1,  Z2 = 1;
   
   bool isNormalKinematics = false;
   bool isOverRideExNegative = true;
   double maxkb = 200.;
   
   //---- beam
   double KEAmean = 100; // MeV/u
   const int nKEA = 1; 
   double KEAList[nKEA] = {300}; 
   
   double KEAsigma = 0; //KEAmean*0.001; // MeV/u , assume Guassian
   double thetaMean = 0.; // mrad 
   double thetaSigma = 0.; // mrad , assume Guassian due to small angle
   
   int numEvent = 100000;
   
   //---- HELIOS detector geometry
   string heliosDetGeoFile = "";//"detectorGeo_upstream.txt";
   double BField = 4.0; // if not detector, must set B-field, else, this value is not used.
   double BFieldTheta = 0.; // direction of B-field
   double eSigma = 0.0001 ; // detector energy sigma MeV
   double zSigma = 0.1 ; // detector position sigma mm
   
   //---- excitation of Beam 
   int nExA = 1;
   double ExAList[nExA];
   ExAList[0] = 0.000; // MeV
   //ExAList[1] = 1.567;
    
   //---- Separation energy
   string separationFile = "separation_energies.txt";
   
   //---- save root file name
   TString saveFileName = "knockout.root";
   
   //---- Auxiliary setting
   bool isTargetScattering = false;
   bool isDecay = false;
   bool isReDo = false; // redo calculation until detected. 
   
   //---- target
   double density = 0.913; // 0.913 g/cm3
   double targetThickness = 2./2. * 2.2e-4; // 2.2 um = 201 ug/cm2
   string stoppingPowerForA = "208Pb_in_CD2.txt"; // generated by SRIM
   string stoppingPowerForb = "1H_in_CD2.txt";
   string stoppingPowerForB = "209Pb_in_CD2.txt";
   
   //=============================================================
   //=============================================================
   //=============================================================
   //===== Set Reaction
   Knockout reaction;
   int AB = AA-A2, ZB = ZA-Z2;
   int A1 = Aa   , Z1 = Za;
   reaction.SetA(AA,ZA);
   reaction.Seta(Aa,Za);
   reaction.Set2(A2,Z2);
   reaction.SetIncidentEnergyAngle(KEAmean, 0, 0);
   reaction.OverRideExNegative(isOverRideExNegative);
   
   printf("===================================================\n");
   printf("=========== %s ===========\n", reaction.GetReactionName().Data());
   printf("=========== KE: %9.4f +- %5.4f MeV/u, dp/p = %5.2f %% \n", KEAmean, KEAsigma, KEAsigma/KEAmean * 50.);
   printf("======== theta: %9.4f +- %5.4f MeV/u \n", thetaMean, thetaSigma);
   printf("===================================================\n");

   
   //======== Set HELIOS
   printf("############################################## HELIOS configuration\n");   
   HELIOS helios1; // for particle-1
   HELIOS helios2; // for particle-2
   helios1.SetMagneticFieldDirection(BFieldTheta);
   helios2.SetMagneticFieldDirection(BFieldTheta);
   bool sethelios1 = helios1.SetDetectorGeometry(heliosDetGeoFile);
   bool sethelios2 = helios2.SetDetectorGeometry(heliosDetGeoFile);
   if( sethelios1 && sethelios2 ) {
		int mDet = helios1.GetNumberOfDetectorsInSamePos();
		printf("========== energy resol.: %f MeV\n", eSigma);
		printf("=========== pos-Z resol.: %f mm \n", zSigma); 
	}else{
		helios1.SetMagneticField(BField);
		helios2.SetMagneticField(BField);
		printf("======== B-field : %5.2f T, Theta: %5.2f deg\n", BField, BFieldTheta);
	}
	
	//==== Target scattering, only energy loss
   if(isTargetScattering) printf("############################################## Target Scattering\n");
   TargetScattering msA;
   TargetScattering msB;
   TargetScattering msb;

   if(isTargetScattering) printf("======== Target : (thickness : %6.2f um) x (density : %6.2f g/cm3) = %6.2f ug/cm2\n", 
                        targetThickness * 1e+4, 
                        density, 
                        targetThickness * density * 1e+6);  
            
   if( isTargetScattering ){
      msA.LoadStoppingPower(stoppingPowerForA);
      msb.LoadStoppingPower(stoppingPowerForb);
      msB.LoadStoppingPower(stoppingPowerForB);
   }
   
   //======= Decay of particle-B
   Decay decay;
   decay.SetMotherDaugther(AB, ZB, AB-1,ZB); //neutron decay
	
	
   //======= loading excitation energy
   printf("############################################## excitation energies\n");
   vector<double> SpList;
   printf("----- loading separation energies.");
   ifstream file;
   file.open(separationFile.c_str());
   string isotopeName;
   if( file.is_open() ){
      string line;
      int i = 0;
      while( file >> line){
         //printf("%d, %s \n", i,  line.c_str());
         if( line.substr(0,2) == "//" ) continue;
         if( i == 0 ) isotopeName = line; 
         if ( i >= 1 ){
            SpList.push_back(atof(line.c_str()));
         }
         i = i + 1;
      }
      file.close();
      printf("... done.\n");
      printf("========== %s\n", isotopeName.c_str());
      int n = SpList.size();
      for(int i = 0; i < n ; i++){
         if( isDecay ) {
            TLorentzVector temp(0,0,0,0);
            int decayID = decay.CalDecay(temp, SpList[i], 0);
            if( decayID == 1) {
               printf("%d, Sp: %6.2f MeV --> Decay. \n", i, SpList[i]);
            }else{
               printf("%d, Sp: %6.2f MeV\n", i, SpList[i]);
            }
         }else{
            printf("%d, Sp: %6.2f MeV \n", i, SpList[i]);
         }
      }
   }else{
       printf("... fail\n");
       return;
   }

   //====================== build tree
   TFile * saveFile = new TFile(saveFileName, "recreate");
   TTree * tree = new TTree("tree", "tree");
   
   double thetaNN, phiNN;
   double theta1, phi1, T1;
   double theta2, phi2, T2;
   double thetaB, TB;
   double Sp, kb, thetab, phib;
   int SpID;

   double ExA;
   int ExAID;
   double KEA, KEAscattered, theta, phi;
   
   double mB,mb;

   tree->Branch("theta1", &theta1, "theta1/D");
   tree->Branch("phi1", &phi1, "phi1/D");
   tree->Branch("T1", &T1, "T1/D");
   tree->Branch("theta2", &theta2, "theta2/D");
   tree->Branch("phi2", &phi2, "phi2/D");
   tree->Branch("T2", &T2, "T2/D");
   tree->Branch("thetaB", &thetaB, "thetaB/D");
   tree->Branch("TB", &TB, "TB/D");
   tree->Branch("thetaNN", &thetaNN, "thetaNN/D");
   tree->Branch("phiNN", &phiNN, "phiNN/D");
   tree->Branch("Sp", &Sp, "Sp/D");
   tree->Branch("kb", &kb, "kb/D");
   tree->Branch("thetab", &thetab, "thetab/D");
   tree->Branch("phib", &phib, "phib/D");
   
   tree->Branch("SpID", &SpID, "SpID/I");
   
   tree->Branch("ExAID", &ExAID, "ExAID/I");
   tree->Branch("KEA", &KEA, "KEA/D");
   if(isTargetScattering) tree->Branch("KEAscattered", &KEAscattered, "KEAscattered/D");
   tree->Branch("theta", &theta, "theta/D");
   tree->Branch("phi", &phi, "phi/D");

   tree->Branch("mB", &mB, "mB/D");
   tree->Branch("mb", &mb, "mb/D");
   
   tree->Branch("Bfield", &BField, "Bfield/D");
   
   TClonesArray * arr = new TClonesArray("TLorentzVector");
   tree->Branch("fV", arr, 256000);
   arr->BypassStreamer();
   
   TClonesArray * arrN = new TClonesArray("TLorentzVector");
   tree->Branch("fVN", arrN, 256000);
   arrN->BypassStreamer();
   
   TLorentzVector* fourVector = NULL;

   // the output of Helios.CalHit
   double e1, z1, t1, rho1, x1h, y1h, r1h;
   double e2, z2, t2, rho2, x2h, y2h, r2h;
	tree->Branch("e1", &e1, "e1/D");
   tree->Branch("z1", &z1, "z1/D");
   tree->Branch("t1", &t1, "t1/D");
   tree->Branch("x1h", &x1h, "x1h/D"); //at 200 mm downstream
   tree->Branch("y1h", &y1h, "y1h/D");
   tree->Branch("r1h", &r1h, "r1h/D");
   tree->Branch("rho1", &rho1, "rho1/D");
   
   tree->Branch("e2", &e2, "e2/D");
   tree->Branch("z2", &z2, "z2/D");
   tree->Branch("t2", &t2, "t2/D");
   tree->Branch("x2h", &x2h, "x2h/D"); //at 200 mm downstream
   tree->Branch("y2h", &y2h, "y2h/D");
   tree->Branch("r2h", &r2h, "r2h/D");
   tree->Branch("rho2", &rho2, "rho2/D");
   
   double Sp2;
   tree->Branch("Sp2", &Sp2, "Sp2/D");
   
   //different coordinate
   double a0, a1, a2; // in k1-k2 coordinate
   tree->Branch("a0", &a0, "a0/D");
   tree->Branch("a1", &a1, "a1/D");
   tree->Branch("a2", &a2, "a2/D");
   
   double b; // in THREEDEE coordinate
   tree->Branch("b", &b, "b/D");
   
   //========timer
   TBenchmark clock;
   bool shown ;   
   clock.Reset();
   clock.Start("timer");
   shown = false;
   printf("############################################## generating %d events \n", numEvent);
   
   //====================================================== calculate 
   int count = 0;
   for( int i = 0; i < numEvent; i++){
      bool redoFlag = true;
      if( !isReDo ) redoFlag = false;
      do{
      
         //==== Set Ex of A
         ExAID = gRandom->Integer(nExA);
         ExA = ExAList[ExAID];
         reaction.SetExA(ExA);
         
         //==== Set Sp of B
         SpID = gRandom->Integer(SpList.size());
         Sp = SpList[SpID]; 
         
         //==== Set incident beam
         if( KEAsigma == 0 ){
            KEA = KEAmean;
         }else{
            KEA = gRandom->Gaus(KEAmean, KEAsigma);
         }
         if( thetaSigma == 0 ){
            theta = thetaMean;
         }else{
            theta = gRandom->Gaus(thetaMean, thetaSigma);
         }
         
         int rKEA = gRandom->Integer(nKEA);
         KEA = KEAList[rKEA];
         
         /*
         KEA = 300*gRandom->Rndm();
         BField = 4*gRandom->Rndm();
			helios1.SetMagneticField(BField);
			helios2.SetMagneticField(BField);
			*/
         
         reaction.SetIncidentEnergyAngle(KEA, theta, 0.);
         
         
         /*
         //For target scattering
         reaction.CalIncidentChannel(isNormalKinematics); // but only need is PA
         TLorentzVector PA = reaction.GetPA();            
         
         double depth = 0;
         if( isTargetScattering ){
            //==== Target scattering, only energy loss
            depth = targetThickness * gRandom->Rndm();
            msA.SetTarget(density, depth); 
            TLorentzVector PAnew = msA.Scattering(PA);
            KEAscattered = msA.GetKE()/AA;
            reaction.SetIncidentEnergyAngle(KEAscattered, theta, phi);
         }*/
         
         //==== Calculate reaction
         thetaNN = TMath::ACos(2 * gRandom->Rndm() - 1) ; 
         phiNN = TMath::TwoPi() * gRandom->Rndm(); 
         
         kb = maxkb * gRandom->Rndm();
         thetab = TMath::ACos(2 * gRandom->Rndm() - 1);
         phib = TMath::TwoPi() * gRandom->Rndm();

         reaction.SetBSpk(Sp, kb, thetab, phib);
         reaction.CalReactionConstant(isNormalKinematics);
         reaction.Event(thetaNN, phiNN);

         TLorentzVector PA = reaction.GetPA();
         TLorentzVector Pa = reaction.GetPa();
               
         TLorentzVector P1 = reaction.GetP1();
         TLorentzVector P2 = reaction.GetP2();
         TLorentzVector Pb = reaction.GetPb();
         TLorentzVector PB = reaction.GetPB();
         
         /*
         //==== Calculate energy loss of scattered and recoil in target
         if( isTargetScattering ){
            if( Pb.Theta() < TMath::PiOver2() ){
               msb.SetTarget(density, targetThickness - depth);
            }else{
               msb.SetTarget(density, depth);
            }
            Pb = msb.Scattering(Pb);
            TbLoss = msb.GetKELoss();
            msB.SetTarget(density, targetThickness - depth);
            PB = msB.Scattering(PB);
         }else{
            TbLoss = 0;
         }
         
         //======= Decay of particle-B
         if( isDecay){
            int decayID = decay.CalDecay(PB, Sp, 0); // decay to ground state
            if( decayID == 1 ){
               PB = decay.GetDaugther_D();
               decayTheta = decay.GetAngleChange();
            }else{
               decayTheta = TMath::QuietNaN();
            }
         }
         */
                  
         //################################### tree branches
         //===== reaction
         theta1 = P1.Theta() * TMath::RadToDeg();
         theta2 = P2.Theta() * TMath::RadToDeg();
         thetaB = PB.Theta() * TMath::RadToDeg();
      
         T1 = P1.E() - P1.M();
         T2 = P2.E() - P2.M();
         TB = PB.E() - PB.M();
         
         phi1 = P1.Phi() * TMath::RadToDeg();
         phi2 = P2.Phi() * TMath::RadToDeg();
         
         //--------- diff coordinate
         b = TMath::ASin( TMath::Sin(P2.Theta()) * TMath::Sin(P1.Phi() - P2.Phi() )) * TMath::RadToDeg();
         
         TVector3 kA = (PA.Vect()).Unit();
         TVector3 k1 = (P1.Vect()).Unit();
         TVector3 k2 = (P2.Vect()).Unit();
         TVector3 n = (k1.Cross(k2)).Unit();
         TVector3 j = (kA - (kA.Dot(n))*n).Unit();
         
         a0 = TMath::ASin(n.Dot(kA)) * TMath::RadToDeg();
         a1 = TMath::ACos(k1.Dot(j)) * TMath::RadToDeg();
         a2 = TMath::ACos(k2.Dot(j)) * TMath::RadToDeg();
         
         //----------- mass
         mB = PB.M();
         mb = Pb.M();
         
         TVector3 bA = PA.BoostVector();
            
         for(int i = 0; i < 6 ; i++){
            TLorentzVector temp;
            double xyzt[4];
            switch(i){
               case 0: temp = PA; break;
               case 1: temp = Pa; break;
               case 2: temp = P1; break;
               case 3: temp = P2; break;
               case 4: temp = PB; break;
               case 5: temp = Pb; break;
            }
            
            temp.GetXYZT(xyzt);
            
            fourVector = (TLorentzVector*) arr->ConstructedAt(i);
            fourVector->SetXYZT(xyzt[0], xyzt[1], xyzt[2], xyzt[3]);
            
            //into normal kinematic
            temp.Boost(-bA);  
            temp.GetXYZT(xyzt);
            
            fourVector = (TLorentzVector*) arrN->ConstructedAt(i);
            fourVector->SetXYZT(xyzt[0], xyzt[1], xyzt[2], xyzt[3]);
            
         }
         
         
         //==== Helios
         int hit1 = helios1.CalHit(P1, Z1, PB, ZB);
         int hit2 = helios2.CalHit(P2, Z2, PB, ZB);
         
         double recoilZ = 1000;
         
         e1 = helios1.GetEnergy() + gRandom->Gaus(0, eSigma);
         z1 = helios1.GetZ() ; 
         t1 = helios1.GetTime();
         rho1 = helios1.GetRho();
         x1h = helios1.GetXPos(recoilZ);
         y1h = helios1.GetYPos(recoilZ);
         r1h = helios1.GetR(recoilZ);
         
         e2 = helios2.GetEnergy() + gRandom->Gaus(0, eSigma);
         z2 = helios2.GetZ() ; 
         t2 = helios2.GetTime();
         rho2 = helios2.GetRho();
         x2h = helios2.GetXPos(recoilZ);
         y2h = helios2.GetYPos(recoilZ);
         r2h = helios2.GetR(recoilZ);
         
         double gammaA = PA.Gamma();
         double betaA = PA.Beta();
         
         Sp2 = (1-gammaA)* 938.272 - gammaA*(e1+e2) + betaA*gammaA*(P1.Pz() + P2.Pz()) ;
         
         //printf("%f, %f | %f, %f \n", e1, z1, e2, z2);
         
         //change thetaNN into deg
         thetaNN = thetaNN * TMath::RadToDeg();
         
         if( hit1 == 1) {
            count ++;
         }
         
         if( isReDo ){
            if( hit1 == 1) {
               redoFlag = false;
            }else{
               redoFlag = true;
               //printf("%d, %2d, thetaNN : %f, theta : %f, z0: %f \n", i, hit, thetaNN * TMath::RadToDeg(), thetab, helios.GetZ0());
            }
         }else{
            redoFlag = false;
         }
         
         
      }while( redoFlag );
      tree->Fill();
      
      //#################################################################### Timer  
      clock.Stop("timer");
      Double_t time = clock.GetRealTime("timer");
      clock.Start("timer");
      
      if ( !shown ) {
         if (fmod(time, 10) < 1 ){
            printf( "%10d[%2d%%]| %8.2f sec | expect: %5.1f min \n", i, TMath::Nint((i+1)*100./numEvent), time , numEvent*time/(i+1)/60);
            shown = 1;
         }
      }else{
         if (fmod(time, 10) > 9 ){
            shown = 0;
         }
      }
   }
   
   saveFile->Write();
   saveFile->Close();
   
   printf("=============== done. saved as %s. count(hit==1) : %d\n", saveFileName.Data(), count);
   
   gROOT->ProcessLine(".q");
   
   /**/

}