#include "CbmL1MuchFinder.h"

#include "TProfile2D.h"
#include "TCanvas.h"
#include "TLegend.h"
#include "TEllipse.h"
#include "TText.h"
#include "TMarker.h"
#include "TStyle.h"

#include "TTree.h"
#include "TFile.h"
#include "TLorentzVector.h"
#include "TVector3.h"
#include "CbmMuchHit.h"
#include "CbmKF.h"
#include "CbmKFMath.h"
#include "CbmRootManager.h"
#include "CbmKFHit.h"
#include "CbmKFPixelMeasurement.h"
#include "CbmKFMaterial.h"
#include "CbmKFTrackInterface.h"
#include "CbmMuchHit.h"
#include "CbmMuchPoint.h"
#include "CbmStsTrack.h"
#include "CbmStsKFTrackFitter.h"

ClassImp(CbmL1MuchFinder);

CbmL1MuchFinder::CbmL1MuchFinder(const char *name, Int_t iVerbose ):CbmTask(name, iVerbose){
}

CbmL1MuchFinder::~CbmL1MuchFinder(){
}

InitStatus CbmL1MuchFinder::Init(){
  return ReInit();
}
  
InitStatus CbmL1MuchFinder::ReInit(){
  fMuchPoints=(TClonesArray *) CbmRootManager::Instance()->GetObject("MuchPoint");
  fMuchHits =(TClonesArray *) CbmRootManager::Instance()->GetObject("MuchHit");
  fStsTracks =(TClonesArray *) CbmRootManager::Instance()->GetObject("STSTrack");

  fStsFitter.Init();

  return kSUCCESS;
}


void CbmL1MuchFinder::SetParContainers(){
}

void CbmL1MuchFinder::Finish(){
}


class CbmL1MuchHit : public CbmKFHit { 

 public:

  CbmL1MuchHit(){}
  ~CbmL1MuchHit(){}

  CbmKFPixelMeasurement FitPoint;

  void Create( CbmMuchHit *h );
  
  void Filter( CbmKFTrackInterface &track, Bool_t downstream, Double_t &QP0 );

  int index;
  int iStation;
  bool busy;
  CbmMuchPoint *MC_Point;

  //ClassDef(CbmL1MuchHit, 1);
};

void CbmL1MuchHit::Create( CbmMuchHit *h )
{
  FitPoint.x = h->GetX();
  FitPoint.y = h->GetY();
  FitPoint.z = h->GetZ();
  FitPoint.V[0] = h->GetDx()*h->GetDx();
  FitPoint.V[1] = 0;
  FitPoint.V[2] = h->GetDy()*h->GetDy();

  CbmKF *KF = CbmKF::Instance();
  MaterialIndex =  KF->GetMaterialIndex( h->GetDetectorID() );
  
  iStation = KF->MuchStationIDMap[ h->GetDetectorID() ];
  busy   = 0;
}

void CbmL1MuchHit::Filter( CbmKFTrackInterface &track, Bool_t downstream, Double_t &QP0 ){  
  track.Extrapolate( FitPoint.z, &QP0 );
  FitPoint.Filter( track );
}


struct CbmL1MuchTrack : public CbmKFTrackInterface{
 public:
  double *GetTrack(){ return T; }
  double *GetCovMatrix(){ return C; }
  double &GetRefChi2(){ return chi2; }
  int    &GetRefNDF(){ return NDF; }
  double GetMass(){ return 0.1057; }
  bool   IsElectron(){ return 0; }
  int    GetNOfHits(){ return vHits.size(); }
  CbmKFHit *GetHit( int i ){ return vHits[i]; }

  double T[6], C[15], chi2 ;
  int NDF ;
  vector<CbmL1MuchHit* > vHits;
  int NHits;
  bool ok;

  void SetStsTrack( CbmStsTrack* track){
    CbmKFMath::CopyTrackParam2TC( track->GetParamLast(), T, C );
    chi2 = track->GetChi2();
    NDF = track->GetNDF();
  }
};

bool ComparePMuchTracks( CbmL1MuchTrack *p1, CbmL1MuchTrack *p2 ){
  return (p1->NHits>p2->NHits) || (p1->NHits==p2->NHits)&&(p1->chi2/p1->NDF <p2->chi2/p2->NDF) ;
}

void CbmL1MuchFinder::Exec(Option_t * option){

  //const double muMass = 0.1057;
  const char fileName[] = "L1MuchFinder.root";
  
  static TFile *HistoFile=0;
  static TTree *HistoTreePull[5];
  static TH1F	*histP[20], *histRes[20][2], *histResMu[20][2], *histChi2[20];
  static TH1F	*histPDG[20];
  static TH1F *histPull[20][5], *histPullMu[20][5];
  static TProfile *histNTr, *histNHits, *histNMu;
  static bool first_call_murec = 1;

  int MuNStations = CbmKF::Instance()->MuchStationIDMap.size();

  if ( first_call_murec )
    {
      first_call_murec = 0;

      HistoTreePull[0] = new TTree("MuRecPullX", "MuRecPullX");
      HistoTreePull[1] = new TTree("MuRecPullY", "MuRecPullY");
      HistoTreePull[2] = new TTree("MuRecPullTx", "MuRecPullTx");
      HistoTreePull[3] = new TTree("MuRecPullTy", "MuRecPullTy");
      HistoTreePull[4] = new TTree("MuRecPullP", "MuRecPullP");

      for( int ist=0; ist<MuNStations; ist++ )
	{
	  char S1[255], S2[255];
	  sprintf(S1,"Psts%i",ist+1);
	  sprintf(S2,"P of sts tracks at Mu detector %i",ist+1);
	  histP[ist] = new TH1F(S1,S2, 500, 0, 30);
	  
	  sprintf(S1,"ResX%i",ist+1);
	  sprintf(S2,"X resolution of sts track at Mu detector %i (cm)",ist+1);
	  histRes[ist][0] = new TH1F(S1,S2,200,-50,50);
	  sprintf(S1,"ResY%i",ist+1);
	  sprintf(S2,"Y resolution of sts track at Mu detector %i (cm)",ist+1);
	  histRes[ist][1] = new TH1F(S1,S2,200,-50,50);
	  
	  sprintf(S1,"ResMuX%i",ist+1);
	  sprintf(S2,"X resolution of Mu sts track at Mu detector %i (cm)",ist+1);
	  histResMu[ist][0] = new TH1F(S1,S2,200,-50,50);
	  sprintf(S1,"ResMuY%i",ist+1);
	  sprintf(S2,"Y resolution of Mu sts track at Mu detector %i (cm)",ist+1);
	  histResMu[ist][1] = new TH1F(S1,S2,200,-50,50);

	  sprintf(S1,"chi2%i",ist+1);
	  sprintf(S2,"chi2 of sts track at Mu detector %i ",ist+1);
	  histChi2[ist] = new TH1F(S1,S2,1000,0,100);
	  sprintf(S1,"PDG%i",ist+1);
	  sprintf(S2,"PDG of the track at Mu detector %i ",ist+1);
	  histPDG[ist] = new TH1F(S1,S2,7000,-3500,3500);
	  
	  sprintf(S1,"PullX_sts%i",ist+1);
	  sprintf(S2,"Pull X of sts tracks at Mu detector %i",ist+1);
	  histPull[ist][0] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullY_sts%i",ist+1);
	  sprintf(S2,"Pull Y of sts tracks at Mu detector %i",ist+1);
	  histPull[ist][1] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullTx_sts%i",ist+1);
	  sprintf(S2,"Pull Tx of sts tracks at Mu detector %i",ist+1);
	  histPull[ist][2] = new TH1F(S1,S2,100,-10,10);

	  sprintf(S1,"PullTy_sts%i",ist+1);
	  sprintf(S2,"Pull Ty of sts tracks at Mu detector %i",ist+1);
	  histPull[ist][3] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullP_sts%i",ist+1);
	  sprintf(S2,"Pull P of sts tracks at Mu detector %i",ist+1);
	  histPull[ist][4] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullMuX_sts%i",ist+1);
	  sprintf(S2,"Pull X of Mu sts tracks at Mu detector %i",ist+1);
	  histPullMu[ist][0] = new TH1F(S1,S2,100,-10,10);

	  sprintf(S1,"PullMuY_sts%i",ist+1);
	  sprintf(S2,"Pull Y of Mu sts tracks at Mu detector %i",ist+1);
	  histPullMu[ist][1] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullMuTx_sts%i",ist+1);
	  sprintf(S2,"Pull Tx of Mu sts tracks at Mu detector %i",ist+1);
	  histPullMu[ist][2] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullMuTy_sts%i",ist+1);
	  sprintf(S2,"Pull Ty of Mu sts tracks at Mu detector %i",ist+1);
	  histPullMu[ist][3] = new TH1F(S1,S2,100,-10,10);
	  
	  sprintf(S1,"PullMuP_sts%i",ist+1);
	  sprintf(S2,"Pull P of Mu sts tracks at Mu detector %i",ist+1);
	  histPullMu[ist][4] = new TH1F(S1,S2,100,-10,10);	  
	}

      histNTr = new TProfile("NTr","N sts tracks per Mu detector",20,1,20);
      histNMu = new TProfile("NMu","N sts muons/N sts tracks per Mu detector",20,1,20);
      histNHits = new TProfile("Nhits","N Mu hits per Mu detector",20,1,20);
    }
  
  int histnhits[MuNStations], histntr[MuNStations], histnmu[MuNStations];
  
  for( int ist=0; ist<MuNStations; ist++ )
      histnhits[ist] = histntr[ist] = histnmu[ist] = 0;
  
  
  int NHits = fMuchHits->GetEntriesFast();
  vector<CbmL1MuchHit> vMuchHits;
  
  for( int ih = 0; ih < NHits; ih++ )
    {
      CbmMuchHit *h = (CbmMuchHit*) fMuchHits->At(ih);
      CbmL1MuchHit m;
      m.Create(h);
      m.index = ih;
      m.MC_Point = (CbmMuchPoint*) fMuchPoints->At(h->GetRefIndex());	
      histnhits[m.iStation]++;
      vMuchHits.push_back(m);
    }
  
  {
    cout <<"Hits per muon station:";
    for( int ist=0; ist<MuNStations; ist++ ){ 
      cout << " " << histnhits[ist]; 
    }
    cout << endl;
  }

  vector<CbmL1MuchTrack> vTracks;

  int NStsTracks = fStsTracks->GetEntriesFast();
  for( int itr=0; itr<NStsTracks; itr++ ){

    CbmStsTrack *tr = (CbmStsTrack*) fStsTracks->At(itr);   

    if ( tr->GetNHits()<4 ) continue;

    fStsFitter.DoFit( tr, 13 ); // refit with muon mass

    CbmL1MuchTrack tmp;
    tmp.SetStsTrack(tr);
    tmp.NHits = 0;
    tmp.ok = 1;
    vTracks.push_back(tmp);    
  }

  int NTracks = vTracks.size();

  for( int itr=0; itr<NTracks; itr++ ){
    CbmL1MuchTrack &tr = vTracks[itr];
    
    bool stopped = 0;
    for( int ist=0; ist<MuNStations && !stopped; ist++ )
      {
	//double qpold = tr.T[4];
	//double zold = tr.T[5];
	double Zdet = CbmKF::Instance()->vMuchDetectors[ist].ZReference;	
	tr.Extrapolate(Zdet);
	if( fabs(tr.T[4])>100.)
	  {
	    cout << "+++++++++ " 
		 << ", stoped track with momentum " << 1./tr.T[4]<<"/"<<1.0/tr.T[4] << " at station "<<ist<< endl;
	    stopped = 1;
	  }
	if( stopped ) break;
	
	if( fabs(1./tr.T[4])<0.5 )
	  {
	    //cout << "..slow track with p " << 1./tr.T[4]<<": "
	    //<<1./qpold<<"/"<<qpold/tr.T[4]<<"/"<<1./tr.T[4]
	    //<<"; z= "<<zold<<"/"<<tr.T[5]<< " at station "<<ist<< endl;
	  }
	int best_hit = -1;
	double best_chi2 = 1.e8;
	for( int ih=0; ih<NHits; ih++ )
	  {
	    CbmL1MuchHit &h = vMuchHits[ih];
	    if( h.iStation != ist ) continue;
	    double dx = tr.T[0] - h.FitPoint.x;
	    double dy = tr.T[1] - h.FitPoint.y;
	    double c0 = tr.C[0] + h.FitPoint.V[0];
	    double c1 = tr.C[1] + h.FitPoint.V[1];
	    double c2 = tr.C[2] + h.FitPoint.V[2];
	    double chi2 = 0.5*(dx*dx*c0-2*dx*dy*c1+dy*dy*c2)/(c0*c2-c1*c1);	    
	    if( chi2<best_chi2 ){
	      best_chi2 = chi2;
	      best_hit = ih;
	    }
	  }

	if ( best_chi2<=10.*10. )
	  {
	    tr.vHits.push_back(&vMuchHits[best_hit]);
	    tr.NHits++;
	    double qp0 = tr.T[4];
	    vMuchHits[best_hit].Filter(tr, 1, qp0);
	  }
      }
    
    //cout<<", pdg = "<<t_pdg<<", chi="<<sqrt(chi2/ndf)<< " mom " << P << ", STS hits = " << t->StsHits.size() + t->MapsHits.size() << " ip_miss: "<< ip_miss<< endl;
  } // tracks
  
  for( int ist=0; ist<MuNStations; ist++ )
    {
      histNTr->Fill(ist+1,histntr[ist]);
      histNHits->Fill(ist+1,histnhits[ist]);
      if( histntr[ist]>0 ) histNMu->Fill(ist+1,histnmu[ist]*1./histntr[ist]);
    }

  //sort tracks and check for common muon hits
 
  vector<CbmL1MuchTrack*> vpTracks;
  for( int i=0; i<NTracks; i++ ) vpTracks.push_back(&(vTracks[i]));
  
  sort( vpTracks.begin(), vpTracks.end(), ComparePMuchTracks );
  for( int it=0; it<NTracks; it++ ){
    CbmL1MuchTrack &t = *vpTracks[it];
    int nbusy = 0;
    for( int ih=0; ih<t.NHits; ih++ ){
      if( t.vHits[ih]->busy ) nbusy++;
    }
    if( nbusy>2 ){
      t.ok = 0; 
      continue;
    }
    for( int ih=0; ih<t.NHits; ih++ ) t.vHits[ih]->busy=1;
  }
  
  HistoTreePull[0]->Fill();
  HistoTreePull[1]->Fill();
  HistoTreePull[2]->Fill();
  HistoTreePull[3]->Fill();
  HistoTreePull[4]->Fill();


  static int EventCounter = 0;
  bool write = false;

  EventCounter++;

  if ( EventCounter < 100 )
	{//write at every 10th event
    if ( EventCounter % 10 == 0)
      write = true;
  }
  
  if ( EventCounter >= 100 )
	{//write at every 100th event
    if ( EventCounter % 100 == 0)
      write = true;
  }
  
  if (write)
	{
    HistoFile = new TFile(fileName,"RECREATE");
    //hPs->Write();
    HistoFile->Write();
    HistoTreePull[0]->Write();
    //RecoData->Print();
    for( int ist=0; ist<MuNStations; ist++ )
		{
      histP[ist]->Write();
      histRes[ist][0]->Write();
      histRes[ist][1]->Write();
      histChi2[ist]->Write();
      histPDG[ist]->Write();
      histPull[ist][0]->Write();
      histPull[ist][1]->Write();
      histPull[ist][2]->Write();
      histPull[ist][3]->Write();
      histPull[ist][4]->Write();
      histPullMu[ist][0]->Write();
      histPullMu[ist][1]->Write();
      histPullMu[ist][2]->Write();
      histPullMu[ist][3]->Write();
      histPullMu[ist][4]->Write();
    }
    histNTr->Write();
    histNMu->Write();
    histNHits->Write();
    HistoFile->Close();
  }
}