// -------------------------------------------------------------------------
// -----                  CbmStsSimulationQa source file               -----
// -----                  Created 02/02/07  by R. Karabowicz           -----
// -------------------------------------------------------------------------


#include "iostream.h"

#include "TCanvas.h"
#include "TPad.h"
#include "TLegend.h"
#include "TPaveText.h"
#include "TClonesArray.h"
#include "TFile.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TH3F.h"
#include "TList.h"
#include "TVector3.h"

#include "CbmGeoNode.h"
#include "CbmGeoPassivePar.h"
#include "CbmGeoStsPar.h"
#include "CbmGeoTransform.h"
#include "CbmGeoVector.h"
#include "CbmMCTrack.h"
#include "CbmMCPoint.h"
#include "CbmStsPoint.h"
#include "CbmRootManager.h"
#include "CbmRunAna.h"
#include "CbmRuntimeDb.h"
#include "CbmStsSimulationQa.h"
#include "CbmStsTrack.h"
#include "CbmStsTrackMatch.h"



// -----   Default constructor   -------------------------------------------
CbmStsSimulationQa::CbmStsSimulationQa() { 
  fOnlineAnalysis = kFALSE;
}
// -------------------------------------------------------------------------



// -----   Standard constructor   ------------------------------------------
CbmStsSimulationQa::CbmStsSimulationQa(Bool_t visualizeBool, Int_t iVerbose) 
  : CbmTask("STS Simulation QA", iVerbose) {
  fOnlineAnalysis = visualizeBool;
}
// -------------------------------------------------------------------------

  
  
// -----   Destructor   ----------------------------------------------------
CbmStsSimulationQa::~CbmStsSimulationQa() { 

  fHistoList->Delete();
  delete fHistoList;
}
// -------------------------------------------------------------------------



// -----   Public method SetParContainers   --------------------------------
void CbmStsSimulationQa::SetParContainers() {

  // Get Run
  CbmRunAna* run = CbmRunAna::Instance();
  if ( ! run ) {
    cout << "-E- " << GetName() << "::SetParContainers: No CbmRunAna!" 
	 << endl;
    return;
  }

  // Get Runtime Database
  CbmRuntimeDb* runDb = run->GetRuntimeDb();
  if ( ! run ) {
    cout << "-E- " << GetName() << "::SetParContainers: No runtime database!" 
	 << endl;
    return;
  }

  // Get passive geometry parameters
  fPassGeo = (CbmGeoPassivePar*) runDb->getContainer("CbmGeoPassivePar");
  if ( ! fPassGeo ) {
    cout << "-E- " << GetName() << "::SetParContainers: "
	 << "No passive geometry parameters!" << endl;
    return;
  }

  // Get STS geometry parameters
  fStsGeo = (CbmGeoStsPar*) runDb->getContainer("CbmGeoStsPar");
  if ( ! fStsGeo ) {
    cout << "-E- " << GetName() << "::SetParContainers: "
	 << "No STS geometry parameters!" << endl;
    return;
  }

}
// -------------------------------------------------------------------------



// -----   Public method Init   --------------------------------------------
InitStatus CbmStsSimulationQa::Init() {

  cout << "==========================================================="
       << endl;;
  cout << GetName() << ": Initialising..." << endl;

  // Get RootManager
  CbmRootManager* ioman = CbmRootManager::Instance();
  if (! ioman) {
    cout << "-E- " << GetName() << "::Init: "
	 << "RootManager not instantised!" << endl;
    return kFATAL;
  }

  // Get MCTrack array
  fMCTracks = (TClonesArray*) ioman->GetObject("MCTrack");
  if ( ! fMCTracks ) {
    cout << "-E- " << GetName() << "::Init: No MCTrack array!" << endl;
    return kERROR;
  }

  // Get STSPoints array
  fSTSPoints = (TClonesArray*) ioman->GetObject("STSPoint");
  if ( ! fSTSPoints ) {
    cout << "-E- " << GetName() << "::Init: No STSPoint array!" << endl;
    return kERROR;
  }
  
  // Get the geometry of target and STS
  InitStatus geoStatus = GetGeometry();
  if ( geoStatus != kSUCCESS ) {
    cout << "-E- " << GetName() << "::Init: Error in reading geometry!"
	 << endl;
    return geoStatus;
  }
  
  // Create histograms
  CreateHistos();
  Reset();
  if ( fOnlineAnalysis ) {
    fOnlineCanvas = new TCanvas("StsSimOnline","Sts simulation online",10,10,600,900);
    fOnlinePad[0] = new TPad("titlePad",   "Title pad"                ,0.01,0.91,0.99,0.99);
    fOnlinePad[1] = new TPad("momentumPad","Momentum pad"             ,0.01,0.46,0.49,0.89);
    fOnlinePad[2] = new TPad("yptPad",     "y-pt pad"                 ,0.01,0.01,0.49,0.44);
    fOnlinePad[3] = new TPad("pointPad",   "Points per track pad"     ,0.51,0.61,0.99,0.89);
    fOnlinePad[4] = new TPad("stationPad", "Points on 1st station pad",0.51,0.31,0.99,0.59);
    fOnlinePad[5] = new TPad("printoutPad","Print information pad    ",0.51,0.01,0.99,0.29);
    fOnlinePad[1]->SetLogy();
    for ( Int_t ipad = 0 ; ipad < 6 ; ipad++ ) {
      fOnlinePad[ipad]->SetFillColor(0);
      fOnlinePad[ipad]->SetBorderMode(0);
      fOnlinePad[ipad]->Draw();
    }
    
    fOnlinePad[0]->cd();
    TLegend* brp = new TLegend(0.1,0.1,0.9,0.9,"Online STS simulation");
    brp->SetTextAlign(22);
    brp->SetTextSize(0.6);
    brp->SetTextColor(1);
    brp->SetBorderSize(0);
    brp->SetFillColor(0);
    brp->Draw();
    fOnlinePad[0]->Update();
  }

  fhNofStsStations->SetBinContent(1,fNStations);

  // Output
  cout << "   Number of STS stations : " << fNStations << endl;
  if (fActive) cout << "   *****   Task is ACTIVE   *****" << endl;
  cout << "==========================================================="
       << endl << endl;

  return geoStatus;

}
// -------------------------------------------------------------------------



// -----   Public method ReInit   ------------------------------------------
InitStatus CbmStsSimulationQa::ReInit() {

  cout << "==========================================================="
       << endl;;
  cout << GetName() << ": Reinitialising..." << endl;

  // Get the geometry of target and STS
  InitStatus geoStatus = GetGeometry();
  if ( geoStatus != kSUCCESS ) {
    cout << "-E- " << GetName() << "::ReInit: Error in reading geometry!"
	 << endl;
    return geoStatus;
  }

  // Output
  cout << "   Number of STS stations : " << fNStations << endl;
  if (fActive) cout << "   *****   Task is ACTIVE   *****" << endl;
  cout << "==========================================================="
       << endl << endl;

  return geoStatus;

}
// -------------------------------------------------------------------------



// -----   Public method Exec   --------------------------------------------
void CbmStsSimulationQa::Exec(Option_t* opt) {

  Int_t nofMCTracks = fMCTracks->GetEntriesFast();
  Int_t nofSTSPoints = fSTSPoints->GetEntriesFast();

  for ( Int_t itr = 0 ; itr < nofMCTracks ; itr++ ) {
    CbmMCTrack *mctrack= (CbmMCTrack*)fMCTracks->At(itr);
    Int_t pdgCode = mctrack->GetPdgCode();

    // reject funny particles
    if ( ( pdgCode == 10010020) ||
	 ( pdgCode == 10010030) ||
	 ( pdgCode == 50000050) ||
	 ( pdgCode == 50010051) ||
	 ( pdgCode == 10020040) )
      continue;

    TVector3 mom = mctrack->GetMomentum();
    Float_t pT  = mom.Pt();
    Float_t p   = mom.Mag();
    
    TVector3 startvtx = mctrack->GetStartVertex();
    Float_t vertexZ  = startvtx.z();
    
    TLorentzVector mom4 = mctrack->Get4Momentum();
    Float_t rapidity =  mom4.Rapidity();
    
    Int_t stsPoints = mctrack->GetStsPoints();
    
    if(stsPoints>0 && vertexZ<=100) { 
      fhMomAll      ->Fill(p);
      fhYPtMapAll   ->Fill(rapidity,pT);
      fhPdgCodeAll  ->Fill(pdgCode);
      fhStsPointsAll->Fill(stsPoints);
      fhMomStsPoints->Fill(stsPoints,p);
    }
    if(stsPoints>3 && vertexZ<=100) { 
      fhMomRec      ->Fill(p);
      fhYPtMapRec   ->Fill(rapidity,pT);
      fhPdgCodeRec  ->Fill(pdgCode);
      fhStsPointsRec->Fill(stsPoints);
    }
  }

  for ( Int_t ipnt = 0 ; ipnt < nofSTSPoints ; ipnt++ ) {
    CbmStsPoint *stsPoint= (CbmStsPoint*)fSTSPoints->At(ipnt);
    Float_t z  = stsPoint->GetZ(); // [cm] 
    Float_t x  = stsPoint->GetX(z); // [cm] 
    Float_t y  = stsPoint->GetY(z); // [cm] 
    
    fhStsPointsPosition->Fill(z,x,y);
    fhStationPoints[fStationNrFromMcId[stsPoint->GetDetectorID()]]->Fill(x,y);
  }
    
  Double_t tracksPerEvent = (Double_t)(fhMomAll->GetEntries())/((Double_t)fNEvents+1);
  Double_t pointsPerEvent = (Double_t)(fhStsPointsPosition->GetEntries())/((Double_t)fNEvents+1);

  if ( fOnlineAnalysis ) {
    fOnlinePad[1]->cd();
    fhMomRec->Draw();
    fOnlinePad[1]->Update();
    fOnlinePad[2]->cd();
    fhYPtMapRec->Draw("colz");
    fOnlinePad[2]->Update();
    fOnlinePad[3]->cd();
    fhStsPointsRec->Draw();
    fOnlinePad[3]->Update();

    fOnlinePad[4]->cd();
    fhStationPoints[0]->Draw();
    fhStationPoints[0]->Draw("colz");
    fOnlinePad[4]->Update();
    
    fOnlinePad[5]->cd();
    TPaveText* printoutPave = new TPaveText(0.1,0.1,0.9,0.9);
    printoutPave->SetTextAlign(22);
    printoutPave->SetTextSize(0.1);
    printoutPave->SetTextColor(1);
    printoutPave->SetBorderSize(0);
    printoutPave->SetFillColor(0);
    printoutPave->AddText(Form("%i events",fNEvents+1));
    printoutPave->AddText(Form("tracks/event = %3.2f",tracksPerEvent));
    printoutPave->AddText(Form("points/event = %3.2f",pointsPerEvent));
    fOnlinePad[5]->Clear();
    printoutPave->Draw();
    fOnlinePad[5]->Update();
  }

  //  cout << "\rEvent #" << fNEvents+1 << flush;
  cout << "\rEvent #" << fNEvents+1 << " --> " 
       << setprecision(6) << tracksPerEvent << " tracks/event >--< " 
       << setprecision(7) << pointsPerEvent << " points/event >--" << flush;

  fNEvents++;

  fhNofEvents->SetBinContent(1,fNEvents);
}
// -------------------------------------------------------------------------



// -----   Private method Finish   -----------------------------------------
void CbmStsSimulationQa::Finish() 
{
  // Run summary to screen
  cout << endl << endl;
  cout << "=======================================================" 
       << endl;
  cout << "            StsSimulationQa: Run summary" << endl << endl;
  cout << "=======================================================" 
       << endl;
  cout << endl << endl;

  gDirectory->mkdir("STSSimulationQA");
  gDirectory->cd("STSSimulationQA");
  TIter next(fHistoList);
  while ( TH1* histo = ((TH1*)next()) ) histo->Write();
  gDirectory->cd("..");
}
// -------------------------------------------------------------------------



// -----   Private method GetGeometry   ------------------------------------
InitStatus CbmStsSimulationQa::GetGeometry() {

  // Get STS geometry
  if ( ! fStsGeo ) {
    cout << "-W- " << GetName() << "::GetGeometry: No passive geometry!"
	 <<endl;
    fNStations = 0;
    return kERROR;
  }
  TObjArray* stsNodes = fStsGeo->GetGeoSensitiveNodes();
  if ( ! stsNodes ) {
    cout << "-E- " << GetName() << "::GetGeometry: No STS node array" 
	 << endl;
    fNStations = 0;
    return kERROR;
  }
  fNStations = stsNodes->GetEntries();

  for ( Int_t ist = 0 ; ist < fNStations ; ist++ ) {
    CbmGeoNode* stsNode = (CbmGeoNode*)stsNodes->At(ist);
    if ( ! stsNode ) {
      cout << "-W- CbmStsDigiScheme::Init: station#" << ist
	   << " not found among sensitive nodes " << endl;
      continue;
    }
    TArrayD* params = stsNode->getParameters();

    fStationsMCId[ist] = stsNode->getMCid();
    fStationNrFromMcId[stsNode->getMCid()] = ist;
    fStationRadius[ist] = params->At(1);

    cout << "station #" << ist+1 << " has MCID = " << stsNode->getMCid() << endl;

  }

  return kSUCCESS;

}
// -------------------------------------------------------------------------



// -----   Private method CreateHistos   -----------------------------------
void CbmStsSimulationQa::CreateHistos() {

  fHistoList = new TList();

  fhMomAll            = new TH1F("hMomAll","Momentum - all in STS",100,0,50);
  fhMomAll->SetXTitle("p [GeV/c]");  fhMomAll->SetYTitle("yield [a.u.]");
  fhYPtMapAll         = new TH2F("hYPtMapAll","Rapidity - trans. mom. map - all in STS",100,-3,7,100,0,5);
  fhYPtMapAll->SetXTitle("rapidity");  fhYPtMapAll->SetYTitle("p_{t} [GeV/c]");
  fhPdgCodeAll        = new TH1F("hPdgCodeAll","PDG code - all in STS",1000,-500,500);
  fhStsPointsAll      = new TH1F("hStsPointsAll","STSPoints per track - all in STS",2*fNStations,0.5,2*fNStations+0.5);
  fhStsPointsAll->SetXTitle("nof points");  fhStsPointsAll->SetYTitle("yield [a.u.]");
  fhMomRec            = new TH1F("hMomRec","Momentum - rec in STS",100,0,50);
  fhMomRec->SetXTitle("p [GeV/c]");  fhMomRec->SetYTitle("yield [a.u.]");
  fhYPtMapRec         = new TH2F("hYPtMapRec","Rapidity - trans. mom. map - rec in STS",100,-3,7,100,0,5);
  fhYPtMapRec->SetXTitle("rapidity");  fhYPtMapRec->SetYTitle("p_{t} [GeV/c]");
  fhPdgCodeRec        = new TH1F("hPdgCodeRec","PDG code - rec in STS",1000,-500,500);
  fhStsPointsRec      = new TH1F("hStsPointsRec","STSPoints per track - rec in STS",2*fNStations,0.5,2*fNStations+0.5);
  fhStsPointsRec->SetXTitle("nof points");  fhStsPointsRec->SetYTitle("yield [a.u.]");
  fhMomStsPoints      = new TH2F("hMomStsPoints","momentum vs STSPoints per track",10,0,9,100,0,50);
  fhStsPointsPosition = new TH3F("hStsPointsPosition","STS hits",100,0,100,100,-50,50,100,-50,50);
  fHistoList->Add(fhMomAll);
  fHistoList->Add(fhYPtMapAll);
  fHistoList->Add(fhPdgCodeAll);
  fHistoList->Add(fhStsPointsAll);
  fHistoList->Add(fhMomRec);
  fHistoList->Add(fhYPtMapRec);
  fHistoList->Add(fhPdgCodeRec);
  fHistoList->Add(fhStsPointsRec);
  fHistoList->Add(fhMomStsPoints);
  fHistoList->Add(fhStsPointsPosition);

  for ( Int_t ist = 0 ; ist < fNStations ; ist++ ) {
    fhStationPoints[ist] = new TH2F(Form("hStationPoints%i",ist+1),
				    Form("Points at station %i",ist+1),
				    1000,-1.1*fStationRadius[ist],1.1*fStationRadius[ist],1000,-1.1*fStationRadius[ist],1.1*fStationRadius[ist]);
    fhStationPoints[ist]->SetXTitle("x [cm]");  fhStationPoints[ist]->SetYTitle("y [cm]");
    fHistoList->Add(fhStationPoints[ist]);
  }

  fhNofEvents      = new TH1F("hNofEvents","Number of events",1,0.,1.);
  fhNofStsStations = new TH1F("hNofStsStations","Number of stations",1,0.,1.);
  fHistoList->Add(fhNofEvents);
  fHistoList->Add(fhNofStsStations);
}
// -------------------------------------------------------------------------



// -----   Private method Reset   ------------------------------------------
void CbmStsSimulationQa::Reset() {
  fNEvents = 0;
}
// -------------------------------------------------------------------------


ClassImp(CbmStsSimulationQa)