//  Geant4 simulation for MuSR
//  AUTHOR: Toni SHIROKA, Paul Scherrer Institut, PSI
//  DATE  : 2008-05
//

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

#include "lem4RootOutput.hh"
#include "G4RunManager.hh"
#include "G4Run.hh"
//#include "globals.hh"
#include "lem4ErrorMessage.hh"

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

lem4RootOutput::lem4RootOutput() {
  pointerToRoot=this;
  boolIsAnySpecialSaveVolumeDefined=false;
  nFieldNomVal=0;
 
  //  SensDetectorMapping["log_ScintBU"]=1;
  //  SensDetectorMapping["log_ScintBD"]=2;
  //  SensDetectorMapping["log_ScintBL"]=3;
  //  SensDetectorMapping["log_ScintBR"]=4;
  //  SensDetectorMapping["log_ScintFU"]=5;
  //  SensDetectorMapping["log_ScintFD"]=6;
  //  SensDetectorMapping["log_ScintFL"]=7;
  //  SensDetectorMapping["log_ScintFR"]=8;
  
  //  SensDetectorMapping["logicScintBU"]=1;
  //  SensDetectorMapping["logicScintBD"]=2;
  //  SensDetectorMapping["logicScintBL"]=3;
  //  SensDetectorMapping["logicScintBR"]=4;
  //  SensDetectorMapping["logicScintFU"]=5;
  //  SensDetectorMapping["logicScintFD"]=6;
  //  SensDetectorMapping["logicScintFL"]=7;
  //  SensDetectorMapping["logicScintFR"]=8;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

lem4RootOutput::~lem4RootOutput() {}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

lem4RootOutput* lem4RootOutput::pointerToRoot=0;
lem4RootOutput* lem4RootOutput::GetRootInstance() {
  return pointerToRoot;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4bool lem4RootOutput::store_runID = true;
G4bool lem4RootOutput::store_eventID = true;
G4bool lem4RootOutput::store_BFieldAtDecay = true;
G4bool lem4RootOutput::store_muIniPosX = true;
G4bool lem4RootOutput::store_muIniPosY = true;
G4bool lem4RootOutput::store_muIniPosZ = true;
G4bool lem4RootOutput::store_muIniMomX = true;
G4bool lem4RootOutput::store_muIniMomY = true;
G4bool lem4RootOutput::store_muIniMomZ = true;
G4bool lem4RootOutput::store_muIniPolX = true;
G4bool lem4RootOutput::store_muIniPolY = true;
G4bool lem4RootOutput::store_muIniPolZ = true;
G4bool lem4RootOutput::store_muDecayDetID= true;
G4bool lem4RootOutput::store_muDecayPosX = true;
G4bool lem4RootOutput::store_muDecayPosY = true;
G4bool lem4RootOutput::store_muDecayPosZ = true;
G4bool lem4RootOutput::store_muDecayTime = true;
G4bool lem4RootOutput::store_muDecayPolX = true;
G4bool lem4RootOutput::store_muDecayPolY = true;
G4bool lem4RootOutput::store_muDecayPolZ = true;
G4bool lem4RootOutput::store_muTargetTime = true;
G4bool lem4RootOutput::store_muTargetPolX = true;
G4bool lem4RootOutput::store_muTargetPolY = true;
G4bool lem4RootOutput::store_muTargetPolZ = true;
G4bool lem4RootOutput::store_posIniMomX = true;
G4bool lem4RootOutput::store_posIniMomY = true;
G4bool lem4RootOutput::store_posIniMomZ = true;
G4bool lem4RootOutput::store_globalTime = true;
G4bool lem4RootOutput::store_fieldValue = true;
G4bool lem4RootOutput::store_det_ID = true;
G4bool lem4RootOutput::store_det_edep = true;
G4bool lem4RootOutput::store_det_edep_el = true;
G4bool lem4RootOutput::store_det_edep_pos = true;
G4bool lem4RootOutput::store_det_edep_gam = true;
G4bool lem4RootOutput::store_det_edep_mup = true;
G4bool lem4RootOutput::store_det_nsteps = true;
G4bool lem4RootOutput::store_det_length = true;
G4bool lem4RootOutput::store_det_start = true;
G4bool lem4RootOutput::store_det_end = true;
G4bool lem4RootOutput::store_det_x = true;
G4bool lem4RootOutput::store_det_y = true;
G4bool lem4RootOutput::store_det_z = true;
G4bool lem4RootOutput::store_fieldNomVal = true;

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void lem4RootOutput::BeginOfRunAction() {
  G4cout << "lem4RootOutput::BeginOfRunAction()  Defining Root tree and branches ..."<<G4endl;
  G4int tmpRunNr=(G4RunManager::GetRunManager())->GetCurrentRun()->GetRunID();
  char RootOutputFileName[200];
  sprintf(RootOutputFileName, "data/lem4_%i.root", tmpRunNr);
  //  G4cout<<"RootOutputFileName="<<RootOutputFileName<<"----"<<G4endl;
  //  rootFile=new TFile("lem4.root","update");
  rootFile=new TFile(RootOutputFileName,"recreate");
  rootTree=new TTree("t1","a simple Tree with simple variables");
  if (store_runID)        {rootTree->Branch("runID",&runID_t,"runID/I");}
  if (store_eventID)      {rootTree->Branch("eventID",&eventID_t,"eventID/I");}
  if (store_BFieldAtDecay){rootTree->Branch("BFieldAtDecay",&B_t,"Bx/D:By:Bz:B3:B4:B5");}
  if (store_muIniPosX)    {rootTree->Branch("muIniPosX",&muIniPosX_t,"muIniPosX/D");}
  if (store_muIniPosY)    {rootTree->Branch("muIniPosY",&muIniPosY_t,"muIniPosY/D");}
  if (store_muIniPosZ)    {rootTree->Branch("muIniPosZ",&muIniPosZ_t,"muIniPosZ/D");}
  if (store_muIniMomX)    {rootTree->Branch("muIniMomX",&muIniMomX_t,"muIniMomX/D");}
  if (store_muIniMomY)    {rootTree->Branch("muIniMomY",&muIniMomY_t,"muIniMomY/D");}
  if (store_muIniMomZ)    {rootTree->Branch("muIniMomZ",&muIniMomZ_t,"muIniMomZ/D");}
  if (store_muIniPolX)    {rootTree->Branch("muIniPolX",&muIniPolX_t,"muIniPolX/D");}
  if (store_muIniPolY)    {rootTree->Branch("muIniPolY",&muIniPolY_t,"muIniPolY/D");}
  if (store_muIniPolZ)    {rootTree->Branch("muIniPolZ",&muIniPolZ_t,"muIniPolZ/D");}
  if (store_muDecayDetID) {rootTree->Branch("muDecayDetID",&muDecayDetID_t,"muDecayDetID/I");}
  if (store_muDecayPosX)  {rootTree->Branch("muDecayPosX",&muDecayPosX_t,"muDecayPosX/D");}
  if (store_muDecayPosY)  {rootTree->Branch("muDecayPosY",&muDecayPosY_t,"muDecayPosY/D");}
  if (store_muDecayPosZ)  {rootTree->Branch("muDecayPosZ",&muDecayPosZ_t,"muDecayPosZ/D");}
  if (store_muDecayTime)  {rootTree->Branch("muDecayTime",&muDecayTime_t,"muDecayTime/D");}
  if (store_muDecayPolX)  {rootTree->Branch("muDecayPolX",&muDecayPolX_t,"muDecayPolX/D");}
  if (store_muDecayPolY)  {rootTree->Branch("muDecayPolY",&muDecayPolY_t,"muDecayPolY/D");}
  if (store_muDecayPolZ)  {rootTree->Branch("muDecayPolZ",&muDecayPolZ_t,"muDecayPolZ/D");}
  if (store_muTargetTime) {rootTree->Branch("muTargetTime",&muTargetTime_t,"muTargetTime/D");}
  if (store_muTargetPolX) {rootTree->Branch("muTargetPolX",&muTargetPolX_t,"muTargetPolX/D");}
  if (store_muTargetPolY) {rootTree->Branch("muTargetPolY",&muTargetPolY_t,"muTargetPolY/D");}
  if (store_muTargetPolZ) {rootTree->Branch("muTargetPolZ",&muTargetPolZ_t,"muTargetPolZ/D");}
  if (store_posIniMomX)   {rootTree->Branch("posIniMomX",&posIniMomx_t,"posIniMomX/D");}
  if (store_posIniMomY)   {rootTree->Branch("posIniMomY",&posIniMomy_t,"posIniMomY/D");}
  if (store_posIniMomZ)   {rootTree->Branch("posIniMomZ",&posIniMomz_t,"posIniMomZ/D");}
  if (store_globalTime)   {rootTree->Branch("globalTime",&globalTime_t,"globalTime/D");}
  if (store_fieldValue)   {rootTree->Branch("fieldValue",&fieldValue_t,"fieldValue/D");}

  //  rootTree->Branch("nSubTracks",&nSubTracks,"nSubTracks/I");
  //  rootTree->Branch("particleID",&particleID_t,"particleID[nSubTracks]/I");
  //  rootTree->Branch("trackID",&trackID_t,"trackID[nSubTracks]/I");
  //  rootTree->Branch("logicalVolumeID",&logicalVolumeID_t,"logicalVolumeID[nSubTracks]/I");
  //  rootTree->Branch("edep",&edep_t,"edep[nSubTracks]/D");
  //  rootTree->Branch("x",&x_t,"x[nSubTracks]/D");
  //  rootTree->Branch("y",&y_t,"y[nSubTracks]/D");
  //  rootTree->Branch("z",&z_t,"z[nSubTracks]/D");
  //  rootTree->Branch("post_x",&post_x_t,"post_x[nSubTracks]/D");
  //  rootTree->Branch("post_y",&post_y_t,"post_y[nSubTracks]/D");
  //  rootTree->Branch("post_z",&post_z_t,"post_z[nSubTracks]/D");
  //  rootTree->Branch("kineticEnergy",&kineticEnergy_t,"kineticEnergy[nSubTracks]/D");
  //  rootTree->Branch("stepLength",&stepLength_t,"stepLength[nSubTracks]/D");
  //  rootTree->Branch("nrOfSubsteps",&nrOfSubsteps_t,"nrOfSubsteps[nSubTracks]/I");

  rootTree->Branch("det_n",&det_n,"det_n/I");
  if (store_det_ID)       {rootTree->Branch("det_ID",&det_ID,"det_ID[det_n]/I");}
  if (store_det_edep)     {rootTree->Branch("det_edep",&det_edep,"det_edep[det_n]/D");}
  if (store_det_edep_el)  {rootTree->Branch("det_edep_el",&det_edep_el,"det_edep_el[det_n]/D");}
  if (store_det_edep_pos) {rootTree->Branch("det_edep_pos",&det_edep_pos,"det_edep_pos[det_n]/D");}
  if (store_det_edep_gam) {rootTree->Branch("det_edep_gam",&det_edep_gam,"det_edep_gam[det_n]/D");}
  if (store_det_edep_mup) {rootTree->Branch("det_edep_mup",&det_edep_mup,"det_edep_mup[det_n]/D");}
  if (store_det_nsteps)   {rootTree->Branch("det_nsteps",&det_nsteps,"det_nsteps[det_n]/I");}
  if (store_det_length)   {rootTree->Branch("det_length",&det_length,"det_length[det_n]/D");}
  if (store_det_start)    {rootTree->Branch("det_time_start",&det_time_start,"det_time_start[det_n]/D");}
  if (store_det_end)      {rootTree->Branch("det_time_end",&det_time_end,"det_time_end[det_n]/D");}
  if (store_det_x)        {rootTree->Branch("det_x",&det_x,"det_x[det_n]/D");}
  if (store_det_y)        {rootTree->Branch("det_y",&det_y,"det_y[det_n]/D");}
  if (store_det_z)        {rootTree->Branch("det_z",&det_z,"det_z[det_n]/D");}

  if (boolIsAnySpecialSaveVolumeDefined) {
    rootTree->Branch("save_n",&save_n,"save_n/I");
    rootTree->Branch("save_detID",&save_detID,"save_detID[save_n]/I");
    rootTree->Branch("save_particleID",&save_particleID,"save_particleID[save_n]/I");
    rootTree->Branch("save_ke",&save_ke,"save_ke[save_n]/D");
    rootTree->Branch("save_x", &save_x, "save_x[save_n]/D");
    rootTree->Branch("save_y", &save_y, "save_y[save_n]/D");
    rootTree->Branch("save_z", &save_z, "save_z[save_n]/D");
    rootTree->Branch("save_px",&save_px,"save_px[save_n]/D");
    rootTree->Branch("save_py",&save_py,"save_py[save_n]/D");
    rootTree->Branch("save_pz",&save_pz,"save_pz[save_n]/D");
  }

  //  htest1 = new TH1F("htest1","The debugging histogram 1",50,-4.,4.);
  //  htest2 = new TH1F("htest2","The debugging histogram 2",50,0.,3.142);
  htest1 = new TH2F("htest1","x, y",50,-200.,200.,50,-200.,200.);
  htest2 = new TH2F("htest2","R, z",50,0.,250.,50,-150.,150.);
  htest3 = new TH1F("htest3","Energy in MeV",55,0.,55.);
  htest4 = new TH1F("htest4","Radioactive electron kinetic energy",250,0.,2.5);
  htest5 = new TH1F("htest5","The debugging histogram 5",50,-4.,4.);
  htest6 = new TH1F("htest6","The debugging histogram 6",50,0.,3.142);
  htest7 = new TH1F("htest7","The debugging histogram 7",50,-4.,4.);
  htest8 = new TH1F("htest8","The debugging histogram 8",50,0.,3.142);

  G4cout << "lem4RootOutput::BeginOfRunAction()  Root tree and branches were defined.\n"<<G4endl;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......


void lem4RootOutput::EndOfRunAction() {
  G4cout << "\nlem4RootOutput::EndOfRunAction() - Writing out the Root tree ..."<<G4endl;
  rootTree->Write();
  htest1->Write();
  htest2->Write();
  htest3->Write();
  htest4->Write();
  htest5->Write();
  htest6->Write();
  htest7->Write();
  htest8->Write();
  // Variables exported from Geant simulation to the Root output
  //  static const Int_t nGeantParamD=10;
  TVectorD TVector_GeantParametersD(maxNGeantParameters);
  for (Int_t i=0; i<maxNGeantParameters; i++) {
    TVector_GeantParametersD[i]=GeantParametersD[i];
  }
  TVector_GeantParametersD.Write("geantParametersD");
  rootFile->Close();
  G4cout<<"lem4RootOutput::EndOfRunAction() - Root tree written out.\n"<<G4endl;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void lem4RootOutput::FillEvent() {
  //  G4cout << "hhhhh: lem4RootOutput::FillEvent()"<<G4endl;
  //  G4cout<<"Kamil: Fill the Root Tree:"<<G4endl;
  //  G4ThreeVector ttt=(posIniMomx_t,posIniMomy_t,posIniMomz_t);
  //  htest5->Fill(atan2(ttt.y(),ttt.x()));
  //  htest6->Fill(atan2(sqrt(ttt.x()*ttt.x()+ttt.y()*ttt.y()),ttt.z())); 
  htest5->Fill(atan2(posIniMomy_t,posIniMomx_t));
  htest6->Fill(atan2(sqrt(posIniMomx_t*posIniMomx_t+posIniMomy_t*posIniMomy_t),posIniMomz_t));
  //  G4double PolModulus=sqrt(muDecayPolX_t*muDecayPolX_t+muDecayPolY_t*muDecayPolY_t+muDecayPolZ_t*muDecayPolZ_t);
  //  G4double PolModulus0=sqrt(muTargetPolX_t*muTargetPolX_t+muTargetPolY_t*muTargetPolY_t+muTargetPolZ_t*muTargetPolZ_t);
  //  if ((muPolx_t>1.0001)||((muPolx_t<-1.0001)&&(muPolx_t>-999.))||(muPolx0_t>1.0001)||((muPolx0_t<-1.0001)&&(muPolx0_t>-999.))) {
  //  if ((PolModulus>1.001)||((PolModulus0>1.001)&&(muTargetTime_t>-999))) {
  //    G4cout <<"lem4RootOutput.cc: Strange polarisation in run="<<runID_t<<" event="<<eventID_t
  //	   <<" PolModulus="<<PolModulus<<" PolModulus0="<<PolModulus0<<G4endl;
  //    G4cout <<"    muPol=("<<muDecayPolX_t<<","<<muDecayPolY_t<<","<<muDecayPolZ_t<<");"
  //           <<"    mupol0=("<<muPolx0_t<<","<<muPoly0_t<<","<<muPolz0_t<<");"<<G4endl;
  //  }
  //  G4cout << "eventID="<<eventID_t<<";  Deposited E (el, pos, gam, mu)="<<det_edep_el<<", "
  //	 << det_edep_pos<<", "<< det_edep_gam<<", "<< det_edep_mup<<G4endl;
  //  G4cout<<"fieldValue="<<fieldValue_t<<G4endl;
  //  G4cout<<"  muDecayPosX="<<muDecayPosX_t << "   muIniPosX="<<muIniPosX_t<<G4endl;
  rootTree->Fill();
  //  G4cout<<"Kamil: end of Filling the Root Tree:"<<G4endl;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void lem4RootOutput::ClearAllRootVariables() {
  runID_t=-1000;
  eventID_t=-1000;
  B_t[0]=-1000.;B_t[1]=-1000.;B_t[2]=-1000.;B_t[3]=-1000.;B_t[4]=-1000.;B_t[5]=-1000.;
  muIniPosX_t=-1000; muIniPosY_t=-1000; muIniPosZ_t=-1000;
  muIniMomX_t=-1000; muIniMomY_t=-1000; muIniMomZ_t=-1000;
  muIniPolX_t=-1000; muIniPolY_t=-1000; muIniPolZ_t=-1000;
  muDecayDetID_t=-1000;
  muDecayPolX_t=-1000; muDecayPolY_t=-1000; muDecayPolZ_t=-1000;
  muTargetPolX_t=-1000; muTargetPolY_t=-1000; muTargetPolZ_t=-1000;
  muDecayPosX_t=-1000;muDecayPosY_t=-1000;muDecayPosZ_t=-1000;
  muDecayTime_t=-1000;
  posIniMomx_t=-1000;posIniMomy_t=-1000;posIniMomz_t=-1000;
  globalTime_t=-1000;
  fieldValue_t=-1000;
  muTargetTime_t=-1000;

  //  nSubTracks=0;
  //  for (G4int i=0; i<maxNSubTracks; i++) {
  //    particleID_t[i]=-1000;
  //    trackID_t[i]=-1000;
  //    logicalVolumeID_t[i]=-1000;
  //    edep_t[i]=-1000;
  //    x_t[i]=-1000;
  //    y_t[i]=-1000;
  //    z_t[i]=-1000;
  //    post_x_t[i]=-1000;
  //    post_y_t[i]=-1000;
  //    post_z_t[i]=-1000;
  //    kineticEnergy_t[i]=-1000;
  //    stepLength_t[i]=-1000;
  //    nrOfSubsteps_t[i]=-1000;
  //  }
  
  det_n=0;
  //  for (G4int i=0; i<NrOfLogicalVolumes; i++) {
  //    det_ID[i]=-1000;
  //    det_edep[i]=-1000;
  //    det_nsteps[i]=-1000;
  //    det_length[i]=-1000;
  //  }

  save_n=0;

}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......


void lem4RootOutput::SetVolumeIDMapping(std::string logivol, int volumeID) {
  //  This function assigns a unique number to each sensitive detector name.
  //  The numbers are used in the root tree, as it is easier to work with numbers
  //  rather than with strings.
  if (SensDetectorMapping[logivol]) {
    char message[100];  
    sprintf(message,"lem4RootOutput::SetVolumeIDMapping: Sensitive volume %s already assigned.",logivol.c_str());
    lem4ErrorMessage::GetInstance()->lem4Error(FATAL,message,false);
  }
  else{
    SensDetectorMapping[logivol]=volumeID;
  }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

G4int lem4RootOutput::ConvertVolumeToID(std::string logivol) {
  G4int volumeID = SensDetectorMapping[logivol];
  if (volumeID==0) {
    char message[100];  
    sprintf(message,"lem4RootOutput::ConvertVolumeToID: No ID number assigned to sensitive volume %s .",logivol.c_str());
    lem4ErrorMessage::GetInstance()->lem4Error(SERIOUS,"message",true);
  }
  return volumeID;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......

void lem4RootOutput::SetSaveDetectorInfo (G4int ID, G4int particleID,  G4double ke, 
                          G4double x, G4double y, G4double z, G4double px, G4double py, G4double pz) {
  if (save_n>=save_nMax) {
    char message[200];
    sprintf(message,"lem4RootOutput.cc::SetSaveDetectorInfo(): more \"save\" hits then allowed: save_nMax=%i",save_nMax);
    lem4ErrorMessage::GetInstance()->lem4Error(SERIOUS,message,true);
  }
  else {
    save_detID[save_n]=ID;
    save_particleID[save_n]=particleID;
    save_ke[save_n]=ke/keV;
    save_x[save_n]=x/mm;
    save_y[save_n]=y/mm;
    save_z[save_n]=z/mm;
    save_px[save_n]=px/MeV;
    save_py[save_n]=py/MeV;
    save_pz[save_n]=pz/MeV;

    save_n++;
  }
}

void lem4RootOutput::SetFieldNomVal(G4int i, G4double value) {
  if (i<maxNFieldnNominalValues) {
    fieldNomVal[i]=value/tesla; /// ATTENTION CHANGE BEFORE, TO ACCOUNT FOR EM FIELDS !! TS
  } 
  else {
    char message[200];
    sprintf(message,
	    "lem4RootOutput.cc::SetFieldNomVal(): more electromagnetic fields then allowed: maxNFieldnNominalValues=%i",
	    maxNFieldnNominalValues);
    lem4ErrorMessage::GetInstance()->lem4Error(SERIOUS,message,true);
  }
}