// -------------------------------------------------------------------------
// -----                CbmLitMuchAna source file            -----
// -----                  Created 01/10/07  by A. Lebedev              -----
// -------------------------------------------------------------------------

#include "CbmLitMuchAna.h"

#include "base/CbmTrackUpdate.h"
#include "base/CbmTrackExtrapolator.h"
#include "base/CbmTrackPropagator.h"
#include "CbmLitKalmanFilter.h"
#include "CbmLitTrackPropagator.h"
#include "CbmLitLineTrackExtrapolator.h"
#include "CbmLitRK4TrackExtrapolator.h"
#include "CbmLitParabolicTrackExtrapolator.h"
//#include "base/CbmTrackPropagatorGeane.h"

#include "CbmMuchHit.h"
#include "CbmMuchTrack.h"
#include "CbmMuchDigi.h"
#include "CbmMuchDigiMatch.h"

// CBM includes
#include "CbmDetectorList.h"
#include "CbmMCPoint.h"
#include "CbmMCTrack.h"
#include "CbmRootManager.h"
#include "CbmRunAna.h"
#include "CbmRuntimeDb.h"

// ROOT includes
#include "TClonesArray.h"
#include "TH1F.h"

// C++ includes
#include <iostream>
#include <iomanip>
#include <map>
#include <cmath>

// -----   Default constructor   -------------------------------------------
CbmLitMuchAna::CbmLitMuchAna() { 
  fMCTrackArray = NULL;
  fMCPointArray = NULL;
  fVerbose      = 1;
  
  fNofTracks15Hits = 0;
  fNofPairTracks15Hits = 0;
  fEvents = 0;
}
// -------------------------------------------------------------------------




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

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



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

   // Get and check CbmRootManager
   CbmRootManager* ioman = CbmRootManager::Instance();
   if ( ! ioman ) Fatal("Init", "No CbmRootManager");
   
   // Get MCTrack array
   fMCTrackArray  = (TClonesArray*) ioman->GetObject("MCTrack");
   if ( ! fMCTrackArray ) Fatal("Init", "No MCTrack array!");
   
   // Get MCPoint array
   fMCPointArray  = (TClonesArray*) ioman->GetObject("MuchPoint");
   if ( ! fMCPointArray ) Fatal("Init", "No MuchPoint array!");
   
   fMuchHits = (TClonesArray*) ioman->GetObject("MuchHit");
   if ( ! fMuchHits ) Fatal("Init", "No fMuchHits array!");
   
      // Get Much digi array
   fMuchDigis  = (TClonesArray*) ioman->GetObject("MuchDigi");
   if ( ! fMuchDigis ) Fatal("Init", "No MuchDigis array!");
   
   // Get Much digi match array
   fMuchDigiMatches  = (TClonesArray*) ioman->GetObject("MuchDigiMatch");
   if ( ! fMuchDigiMatches ) Fatal("Init", "No MuchDigiMatches array!");
      
   fMuchTracks = (TClonesArray*) ioman->GetObject("MuchTrack");
   if ( ! fMuchTracks ) Fatal("Init", "No trackArray!");
     
     
   fExtrapolator = new CbmLitRK4TrackExtrapolator();
   fExtrapolator = new CbmLitParabolicTrackExtrapolator();
   fExtrapolator->Initialize();
   fPropagator = new CbmLitTrackPropagator(fExtrapolator);
   //fPropagator = new CbmTrackPropagatorGeane();
   if(NULL == fPropagator) {
      std::cout << "-E- CbmLitMuchAna::Init : "
           << "propagator is not initialized." << std::endl;
   }
   fPropagator->Properties().SetProperty("fMass",0.105);
   fPropagator->Properties().SetProperty("fApplyEnergyLoss", true);
   fPropagator->Properties().SetProperty("fEnergyLoss", 0.00354);
   fPropagator->Properties().SetProperty("fFms", 1.25);
   fPropagator->Initialize();

  
   fFilter = new CbmLitKalmanFilter();
   if(NULL == fFilter) {
      std::cout << "-E- CbmLitMuchAna::Init : "
           << "filter is not initialized." << std::endl;
   }
   fFilter->Initialize();
   
   
   fNofLayers = 10;
   
   fh_srx.resize(fNofLayers);
   fh_sry.resize(fNofLayers);
   
   fh_resx.resize(fNofLayers);
   fh_resy.resize(fNofLayers);
   fh_restx.resize(fNofLayers);
   fh_resty.resize(fNofLayers);
   fh_resqp.resize(fNofLayers);
   
   fh_pullx.resize(fNofLayers);
   fh_pully.resize(fNofLayers);
   fh_pulltx.resize(fNofLayers);
   fh_pullty.resize(fNofLayers);
   fh_pullqp.resize(fNofLayers);
   
   fh_resp.resize(fNofLayers);
   
   char histName[100];
   char histTitle[100];
   
   for (Int_t i = 0; i < fNofLayers; i++) {
       
      sprintf(histName, "fh_srx%i",i);
      sprintf(histTitle, "Search Region X. X_pred - X_hit at %i MUCH layer",i);
      fh_srx[i] = new TH1F(histName, histTitle, 200, -15., 15.);
                                   
      sprintf(histName, "fh_sry%i",i);
      sprintf(histTitle, "Search Region Y. Y_pred - Y_hit at %i MUCH layer",i);
      fh_sry[i] = new TH1F(histName, histTitle, 200, -15., 15.); 
   
   
      sprintf(histName, "fh_resx%i",i);
      sprintf(histTitle, "Resolution X. X_pred - X_point at %i MUCH layer",i);
      fh_resx[i] = new TH1F(histName, histTitle, 200, -15., 15.);
                                   
      sprintf(histName, "fh_resy%i",i);
      sprintf(histTitle, "Resolution Y. Y_pred - Y_point at %i MUCH layer",i);
      fh_resy[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      sprintf(histName, "fh_restx%i",i);
      sprintf(histTitle, "Resolution Tx. Tx_pred - Px/Pz at %i MUCH layer",i);
      fh_restx[i] = new TH1F(histName, histTitle, 200, -15., 15.);
                                   
      sprintf(histName, "fh_resty%i",i);
      sprintf(histTitle, "Resolution Ty. Ty_pred - Py/Pz at %i MUCH layer",i);
      fh_resty[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      sprintf(histName, "fh_resqp%i",i);
      sprintf(histTitle, "Resolution Q/P. qp_pred - qp_point at %i MUCH layer",i);
      fh_resqp[i] = new TH1F(histName, histTitle, 200, -15., 15.);      
      
            
      sprintf(histName, "fh_pullx%i",i);
      sprintf(histTitle, "Pull X at%i MUCH layer",i);
      fh_pullx[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      sprintf(histName, "fh_pully%i",i);
      sprintf(histTitle, "Pull Y at %i MUCH layer",i);
      fh_pully[i] = new TH1F(histName, histTitle, 200, -15., 15.);
       
      sprintf(histName, "fh_pulltx%i",i);
      sprintf(histTitle, "Pull Tx at %i MUCH layer",i);
      fh_pulltx[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      sprintf(histName, "fh_pullty%i",i);
      sprintf(histTitle, "Pull Ty at %i MUCH layer",i);
      fh_pullty[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      sprintf(histName, "fh_pullqp%i",i);
      sprintf(histTitle, "Pull q/p at %i MUCH layer",i);
      fh_pullqp[i] = new TH1F(histName, histTitle, 200, -15., 15.);
      
      
      sprintf(histName, "fh_resp%i",i);
      sprintf(histTitle, "Res % p at %i MUCH layer",i);
      fh_resp[i] = new TH1F(histName, histTitle, 200, -15., 15.);
   }
   
}
// -------------------------------------------------------------------------

// -----  SetParContainers -------------------------------------------------
void CbmLitMuchAna::SetParContainers()
{
    CbmRunAna* ana = CbmRunAna::Instance();
    CbmRuntimeDb* rtdb = ana->GetRuntimeDb();

    rtdb->getContainer("CbmBaseParSet");
    rtdb->getContainer("CbmGeoPassivePar");
    rtdb->getContainer("CbmGeoStsPar");
    rtdb->getContainer("CbmGeoRichPar");
    rtdb->getContainer("CbmGeoMuchPar");
    rtdb->getContainer("CbmGeoTrdPar");
    rtdb->getContainer("CbmFieldPar");
}

// -----   Public method DoFind   ------------------------------------------
void CbmLitMuchAna::Exec(Option_t* opt) 
{
  std::cout << "---------CbmLitMuchAna::Exec-----------------------" << std::endl;
  
  Int_t nofTracks = fMuchTracks->GetEntriesFast();
  //std::cout << "Number of MUCH tracks: " << nofTracks << std::endl;
  
  for (Int_t iTrack = 0; iTrack < nofTracks; iTrack++) {
  
      CbmMuchTrack* pTrack = (CbmMuchTrack*) fMuchTracks->At(iTrack);
      
      Int_t nofHits = pTrack->GetNHits();
      if(nofHits != fNofLayers) continue;
      
      std::vector<CbmMuchHit*> pHits;
      std::vector<CbmMCPoint*> pPoints;
      
      for (Int_t iHit = 0; iHit < nofHits; iHit++){
      
         Int_t hitIndex = pTrack->GetHitIndex(iHit);
         CbmMuchHit* pHit = (CbmMuchHit*) fMuchHits->At(hitIndex);
         pHits.push_back(pHit);
         
         Int_t digiIndex = pHit->GetDigi();
         CbmMuchDigiMatch* pDigiMatch = (CbmMuchDigiMatch*) fMuchDigiMatches->At(digiIndex);    
         Int_t pointIndex = pDigiMatch->GetRefIndex(0);
         
         if (pointIndex < 0) Fatal("MUCH Ana", "Wrong PointIndex");
                  
         CbmMCPoint* pPoint = (CbmMCPoint*) fMCPointArray->At(pointIndex);
         
         if (!pPoint) Fatal("MUCH Ana", "Wrong Point pointer");
         
         pPoints.push_back(pPoint);
         
      }
      

      //CbmTrackParam* par;
      //par = pTrack->GetMuchTrack();
      
      CbmTrackParam *par = new CbmTrackParam;
      
      CbmTrackParam *parIn;
      parIn = pTrack->GetMuchTrack();
      
      for (Int_t iHit = 0; iHit < nofHits; iHit++){
                    
         Double_t zOut = pHits[iHit]->GetZ(); 
               
        
         fPropagator->Propagate(parIn, par, zOut);
         
         fh_srx[iHit]->Fill(par->GetX() - pHits[iHit]->GetX());
         fh_sry[iHit]->Fill(par->GetY() - pHits[iHit]->GetY());
         
         
         //fFilter->Update(par, pHits[iHit]);
         
         
         Double_t x = pPoints[iHit]->GetX();
         Double_t y = pPoints[iHit]->GetY();
         Double_t px = pPoints[iHit]->GetPx();
         Double_t py = pPoints[iHit]->GetPy();
         Double_t pz = pPoints[iHit]->GetPz();
         Double_t tx = px / pz;
         Double_t ty = py / pz;
         
         Int_t mcId = pPoints[iHit]->GetTrackID();
         CbmMCTrack* mcTrack = (CbmMCTrack*) fMCTrackArray->At(mcId);
         Int_t pdg = mcTrack->GetPdgCode();
         Double_t q;
         if (pdg > 0) q = -1.; else q = 1.;         
         Double_t qp = q / sqrt(px * px + py * py + pz * pz);
         
         Double_t resx = par->GetX() - x;
         Double_t resy = par->GetY() - y;
         Double_t restx = par->GetTx() - tx;
         Double_t resty = par->GetTy() - ty;
         Double_t resqp = par->GetQp() - qp;         
         
         Double_t resp = 100 * (((1. / fabs(par->GetQp())) - (1. / fabs(qp))) / (1. / qp));   
                  
         Double_t pullx = (par->GetX() - x) / (sqrt(par->GetCovariance(0,0)));
         Double_t pully = (par->GetY() - y) / (sqrt(par->GetCovariance(1,1)));
         Double_t pulltx = (par->GetTx() - tx) / (sqrt(par->GetCovariance(2,2)));
         Double_t pullty = (par->GetTy() - ty) / (sqrt(par->GetCovariance(3,3)));
         Double_t pullqp = (par->GetQp() - qp) / (sqrt(par->GetCovariance(4,4)));
         //std::cout << "GetQp = " << par->GetQp()<< ", qp = " << qp << ", pull qp = " << pullqp << std::endl;
         
         
         fh_resx[iHit]->Fill(resx);
         fh_resy[iHit]->Fill(resy); 
         fh_restx[iHit]->Fill(restx);
         fh_resty[iHit]->Fill(resty);
         fh_resqp[iHit]->Fill(resqp);
         
         fh_pullx[iHit]->Fill(pullx);
         fh_pully[iHit]->Fill(pully); 
         fh_pulltx[iHit]->Fill(pulltx);
         fh_pullty[iHit]->Fill(pullty);
         fh_pullqp[iHit]->Fill(pullqp);
         
         fh_resp[iHit]->Fill(resp);
         
              
         //fFilter->Update(par, pHits[iHit]);
                   
         
      }
     // pTrack->SetMuchTrack(par);
     
     delete par;
                   
  }
  
  
            
  fEvents++;
  std::cout << "Event number: " << fEvents << std::endl;
  std::cout << "---------------------------------------------" << std::endl;

  //return nofTracks;

}
// -------------------------------------------------------------------------
void CbmLitMuchAna::Finish()
{
   for (Int_t i = 0; i < fNofLayers; i++) {
      fh_srx[i]->Write();
      fh_sry[i]->Write();
      
      fh_resy[i]->Write();
      fh_resx[i]->Write();
      fh_resty[i]->Write();
      fh_restx[i]->Write();      
      fh_resqp[i]->Write(); 
      
      fh_pully[i]->Write();
      fh_pullx[i]->Write();
      fh_pullty[i]->Write();
      fh_pulltx[i]->Write();      
      fh_pullqp[i]->Write();     
      
      fh_resp[i]->Write();  
   }
}


ClassImp(CbmLitMuchAna)