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

#include "CbmMvdStripsSimulationQa.h"

#include "CbmMvdStripsPoint.h"
// #include "CbmMvdStripsTrack.h"
//#include "CbmMvdStripsTrackMatch.h"
#include "CbmGeoMvdStripsPar.h"

#include "CbmDetectorList.h"
#include "FairGeoNode.h"
#include "CbmGeoPassivePar.h"
#include "FairGeoTransform.h"
#include "FairGeoVector.h"
#include "CbmMCTrack.h"
#include "FairMCPoint.h"
#include "FairRootManager.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.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 <iostream>
#include <iomanip>

using std::cout;
using std::endl;
using std::flush;
using std::setprecision;


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



// -----   Standard constructor   ------------------------------------------
CbmMvdStripsSimulationQa::CbmMvdStripsSimulationQa(Bool_t visualizeBool, Int_t iVerbose)
  : FairTask("MVDSTRIPS Simulation QA", iVerbose) {
  fOnlineAnalysis = visualizeBool;
}
// -------------------------------------------------------------------------



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

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



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

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

  // Get Runtime Database
  FairRuntimeDb* 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 MVDSTRIPS geometry parameters
  fMvdStripsGeo = (CbmGeoMvdStripsPar*) runDb->getContainer("CbmGeoMvdStripsPar");
  if ( ! fMvdStripsGeo ) {
    cout << "-E- " << GetName() << "::SetParContainers: "
	 << "No MvdStrips geometry parameters!" << endl;
    return;
  }

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



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

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

  // Get RootManager
  FairRootManager* ioman = FairRootManager::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 MvdStripsPoints array
  fMVDSTRIPSPoints = (TClonesArray*) ioman->GetObject("MvdStripsPoint");
  if ( ! fMVDSTRIPSPoints ) {
    cout << "-E- " << GetName() << "::Init: No MvdStripsPoint array!" << endl;
    return kERROR;
  }

  // Get the geometry of target and MvdStrips
  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("MvdStripsSimOnline","MvdStrips simulation online",10,10,600,600);
    fOnlinePad[0] = new TPad("titlePad",      "Title pad"             ,0.00,0.90,1.00,1.00);
    fOnlinePad[1] = new TPad("momentumPad",   "Momentum pad"          ,0.00,0.35,0.50,0.90);
    fOnlinePad[2] = new TPad("printoutPad","Print information pad    ",0.10,0.10,0.35,0.35);
    fOnlinePad[3] = new TPad("pointPad",      "Points per track pad"  ,0.50,0.50,1.00,0.90);
    fOnlinePad[4] = new TPad("stationPad1","Points on 1st station pad",0.50,0.25,0.75,0.50);
    fOnlinePad[5] = new TPad("stationPad2","Points on 3rd station pad",0.75,0.25,1.00,0.50);
    fOnlinePad[6] = new TPad("stationPad3","Points on 5th station pad",0.50,0.00,0.75,0.25);
    fOnlinePad[7] = new TPad("stationPad4","Points on 7th station pad",0.75,0.00,1.00,0.25);
    fOnlinePad[1]->SetLogy();
    for ( Int_t ipad = 0 ; ipad < 8 ; 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 MvdStrips simulation");
    brp->SetTextAlign(22);
    brp->SetTextSize(0.6);
    brp->SetTextColor(1);
    brp->SetBorderSize(0);
    brp->SetFillColor(0);
    brp->Draw();
    fOnlinePad[0]->Update();
  }

  fhNofMvdStripsStations->SetBinContent(1,fNStations);

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

  return geoStatus;

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



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

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

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

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

  return geoStatus;

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



// -----   Public method Exec   --------------------------------------------
void CbmMvdStripsSimulationQa::Exec(Option_t* opt) {
  //  cout << "EXEC" << endl;

  Int_t nofMCTracks = fMCTracks->GetEntriesFast();
  Int_t nofMVDSTRIPSPoints = fMVDSTRIPSPoints->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(mom);
    Float_t pT  = mom.Pt();
    Float_t p   = mom.Mag();

    TVector3 startvtx;
    mctrack->GetStartVertex(startvtx);
    Float_t vertexZ  = startvtx.z();

    TLorentzVector mom4;
    mctrack->Get4Momentum(mom4);
    Float_t rapidity =  mom4.Rapidity();

    Int_t mvdstripsPoints = mctrack->GetNPoints(kMVDSTRIPS);

    if(mvdstripsPoints>0 && vertexZ<=100) {
      fhMomAll      ->Fill(p);
      fhYPtMapAll   ->Fill(rapidity,pT);
      fhPdgCodeAll  ->Fill(pdgCode);
      fhMvdStripsPointsAll->Fill(mvdstripsPoints);
      fhMomMvdStripsPoints->Fill(mvdstripsPoints,p);
    }
    if(mvdstripsPoints>3 && vertexZ<=100) {
      fhMomRec      ->Fill(p);
      fhYPtMapRec   ->Fill(rapidity,pT);
      fhPdgCodeRec  ->Fill(pdgCode);
      fhMvdStripsPointsRec->Fill(mvdstripsPoints);
    }
  }

  //  cout << "track loop done" << endl;

  for ( Int_t ipnt = 0 ; ipnt < nofMVDSTRIPSPoints ; ipnt++ ) {
    CbmMvdStripsPoint *mvdstripsPoint= (CbmMvdStripsPoint*)fMVDSTRIPSPoints->At(ipnt);
    Float_t z  = mvdstripsPoint->GetZ(); // [cm]
    Float_t x  = mvdstripsPoint->GetX(z); // [cm]
    Float_t y  = mvdstripsPoint->GetY(z); // [cm]

    fhMvdStripsPointsPosition->Fill(z,x,y);
//     cout << "filled 3d,    MCId = " << mvdstripsPoint->GetDetectorID()
// 	 << "             nr = " << fStationNrFromMcId[mvdstripsPoint->GetDetectorID()] << endl;
    fhStationPoints[fStationNrFromMcId[mvdstripsPoint->GetDetectorID()]]->Fill(x,y);
//     cout << "filled 2d for MCId = " << mvdstripsPoint->GetDetectorID()
// 	 << " and station nr = " << fStationNrFromMcId[mvdstripsPoint->GetDetectorID()] << endl;
  }

  //  cout << "hello" << endl;

  Double_t tracksPerEvent = (Double_t)(fhMomAll->GetEntries())/((Double_t)fNEvents+1);
  Double_t pointsPerEvent = (Double_t)(fhMvdStripsPointsPosition->GetEntries())/((Double_t)fNEvents+1);

  if ( fOnlineAnalysis ) {
    fOnlinePad[1]->cd();
    fhMomRec->Draw();
    fOnlinePad[1]->Update();
    fOnlinePad[2]->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[2]->Clear();
    printoutPave->Draw();
    fOnlinePad[2]->Update();
    fOnlinePad[3]->cd();
    fhMvdStripsPointsRec->Draw();
    fOnlinePad[3]->Update();
 
    if ( fNStations ) {
      fOnlinePad[4]->cd();
      fhStationPoints[0]->Draw("colz");
      fOnlinePad[4]->Update();
      if ( fNStations>2 ) {
	fOnlinePad[5]->cd();
	fhStationPoints[2]->Draw("colz");
	fOnlinePad[5]->Update();
	if ( fNStations>4 ) {
	  fOnlinePad[6]->cd();
	  fhStationPoints[4]->Draw("colz");
	  fOnlinePad[6]->Update();
	  if ( fNStations>6 ) {
	    fOnlinePad[7]->cd();
	    fhStationPoints[6]->Draw("colz");
	    fOnlinePad[7]->Update();
	  }
	}
      }
    }

    
  }

  //  cout << "\rEvent #" << fNEvents+1 << flush;
  cout << endl << endl;
  cout << "======================================================="<< endl;
  cout << "===== MvdStripsSimulationQa: " << endl;
  cout << "===== \rEvent #" << fNEvents+1 << endl;
  cout << "===== " << setprecision(6) << tracksPerEvent << " tracks/event" << endl;
  cout << "===== " << setprecision(7) << pointsPerEvent << " points/event" << endl;
  cout << "======================================================="<< endl;

  fNEvents++;

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



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

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



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

  cout << "GET GEOMETRY" << endl;
  // Get MvdStrips geometry
  if ( ! fMvdStripsGeo ) {
    cout << "-W- " << GetName() << "::GetGeometry: No passive geometry!"
	 <<endl;
    fNStations = 0;
    return kERROR;
  }
  TObjArray* mvdstripsNodes = fMvdStripsGeo->GetGeoSensitiveNodes();
  if ( ! mvdstripsNodes ) {
    cout << "-E- " << GetName() << "::GetGeometry: No MvdStrips node array"
	 << endl;
    fNStations = 0;
    return kERROR;
  }
  Int_t tempNofStations = mvdstripsNodes->GetEntries();

//   cout << "There are " << tempNofStations << " nodes" << (tempNofStations > 10 ? "!!!" : "" ) << endl;

  TString geoNodeName;
  fNStations = 0;
  TString stationNames[10];
  for ( Int_t ist = 0 ; ist < tempNofStations ; ist++ ) {
    FairGeoNode* mvdstripsNode = (FairGeoNode*)mvdstripsNodes->At(ist);
    if ( ! mvdstripsNode ) {
      cout << "-W- CbmMvdStripsDigiScheme::Init: station#" << ist
	   << " not found among sensitive nodes " << endl;
      continue;
    }
    geoNodeName = mvdstripsNode->GetName();
    TArrayD* params = mvdstripsNode->getParameters();

    Bool_t stationKnown = kFALSE;
    // check if the node belongs to some station, save the MCId and outer radius
    for ( Int_t ikst = 0 ; ikst < fNStations ; ikst++ )
      if ( geoNodeName.Contains(stationNames[ikst]) ) {
	FairGeoTransform* transform = mvdstripsNode->getLabTransform();
	FairGeoVector translat = transform->getTranslation();
	FairGeoTransform center = mvdstripsNode->getCenterPosition();
	FairGeoVector centerV = center.getTranslation();
	Double_t radialSize = TMath::Sqrt((translat.X() + centerV.X() + params->At(0))*
					  (translat.X() + centerV.X() + params->At(0)) +
					  (translat.Y() + centerV.Y() + params->At(1))*
					  (translat.Y() + centerV.Y() + params->At(1)));
	if ( radialSize > fStationRadius[ikst] ) {
// 	  cout << "Found bigger radius for station " << ikst << " -> " << radialSize << endl;
	  fStationRadius[ikst] = radialSize;
	}
	fStationNrFromMcId[mvdstripsNode->getMCid()] = ikst;
	stationKnown = kTRUE;
      }

    if ( stationKnown ) continue;

    // if not known, register it and save MCId
    fStationNrFromMcId[mvdstripsNode->getMCid()] = fNStations;

    if ( geoNodeName.Length() == 12 )
      fStationRadius[fNStations] = params->At(1);
    else {
      FairGeoTransform* transform = mvdstripsNode->getLabTransform();
      FairGeoVector translat = transform->getTranslation();
      FairGeoTransform center = mvdstripsNode->getCenterPosition();
      FairGeoVector centerV = center.getTranslation();
      Double_t radialSize = TMath::Sqrt((translat.X() + centerV.X() + params->At(0))*
					(translat.X() + centerV.X() + params->At(0)) +
					(translat.Y() + centerV.Y() + params->At(1))*
					(translat.Y() + centerV.Y() + params->At(1)));
      fStationRadius[fNStations] = radialSize;
    }
    // it will work only if the node name is organized as:
    // for station name is "mvdstripsstationXX", where XX is the station number (f.e. XX=07 for station number 7)
    // for sector  name is "mvdstripsstationXXanythingHereToDistinguishDifferentSectors"
    geoNodeName.Remove(12,geoNodeName.Length()-12);
    stationNames[fNStations] = geoNodeName.Data();
    fNStations++;

//     cout << "station #" << fNStations << " has MCID = " << mvdstripsNode->getMCid() << " and name " << mvdstripsNode->GetName() << endl;

    //    fStationsMCId[fNStations] = mvdstripsNode->getMCid(); // not used
  }
//   cout << "There are " << fNStations << " stations" << endl;

  return kSUCCESS;

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



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

  fHistoList = new TList();

  fhMomAll            = new TH1F("hMomAll","Momentum - all in MVDSTRIPS",100,0,50);
  fhMomAll->SetXTitle("p [GeV/c]");  fhMomAll->SetYTitle("yield [a.u.]");
  fhYPtMapAll         = new TH2F("hYPtMapAll","Rapidity - trans. mom. map - all in MVDSTRIPS",100,-3,7,100,0,5);
  fhYPtMapAll->SetXTitle("rapidity");  fhYPtMapAll->SetYTitle("p_{t} [GeV/c]");
  fhPdgCodeAll        = new TH1F("hPdgCodeAll","PDG code - all in MVDSTRIPS",1000,-500,500);
  fhMvdStripsPointsAll      = new TH1F("hMvdStripsPointsAll","MVDSTRIPSPoints per track - all in MVDSTRIPS",2*fNStations,0.5,2*fNStations+0.5);
  fhMvdStripsPointsAll->SetXTitle("nof points");  fhMvdStripsPointsAll->SetYTitle("yield [a.u.]");
  fhMomRec            = new TH1F("hMomRec","Momentum - rec in MVDSTRIPS",100,0,50);
  fhMomRec->SetXTitle("p [GeV/c]");  fhMomRec->SetYTitle("yield [a.u.]");
  fhYPtMapRec         = new TH2F("hYPtMapRec","Rapidity - trans. mom. map - rec in MVDSTRIPS",100,-3,7,100,0,5);
  fhYPtMapRec->SetXTitle("rapidity");  fhYPtMapRec->SetYTitle("p_{t} [GeV/c]");
  fhPdgCodeRec        = new TH1F("hPdgCodeRec","PDG code - rec in MVDSTRIPS",1000,-500,500);
  fhMvdStripsPointsRec      = new TH1F("hMvdStripsPointsRec","MVDSTRIPSPoints per track - rec in MVDSTRIPS",2*fNStations,0.5,2*fNStations+0.5);
  fhMvdStripsPointsRec->SetXTitle("nof points");  fhMvdStripsPointsRec->SetYTitle("yield [a.u.]");
  fhMomMvdStripsPoints      = new TH2F("hMomMvdStripsPoints","momentum vs MVDSTRIPSPoints per track",1000,0,9,100,0,50);
  fhMvdStripsPointsPosition = new TH3F("hMvdStripsPointsPosition","MVDSTRIPS hits",100,0,100,100,-50,50,100,-50,50);
  fHistoList->Add(fhMomAll);
  fHistoList->Add(fhYPtMapAll);
  fHistoList->Add(fhPdgCodeAll);
  fHistoList->Add(fhMvdStripsPointsAll);
  fHistoList->Add(fhMomRec);
  fHistoList->Add(fhYPtMapRec);
  fHistoList->Add(fhPdgCodeRec);
  fHistoList->Add(fhMvdStripsPointsRec);
  fHistoList->Add(fhMomMvdStripsPoints);
  fHistoList->Add(fhMvdStripsPointsPosition);

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

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



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


ClassImp(CbmMvdStripsSimulationQa)