#include "CbmRichAnalysisRingFinder.h"
#include "CbmRootManager.h"
#include "CbmMCTrack.h"
#include "CbmRichHit.h"
#include "CbmRichPoint.h"
#include "CbmRichRing.h"
#include "CbmRichLightSpot.h"
#include "TH1.h"
#include "TH2.h"
#include "TObjArray.h"

//------------------------------------------------------------------------------
CbmRichAnalysisRingFinder::CbmRichAnalysisRingFinder() :CbmTask()
{
}

//------------------------------------------------------------------------------
CbmRichAnalysisRingFinder::CbmRichAnalysisRingFinder(const char *name, const char *title)
  :CbmTask(name)
{
  //Histograms
  fhRingPabs=new TH1F("hRingPabs","Momentum of track associated with ring",100,0,50);
  fhSpotPabs=new TH1F("hSpotPabs","Momentum of track associated with LighSpot",100,0,50);

  fhElectronRingPabs = 
    new TH1F("hElectronRingPabs","Momentum associated with electron ring",100,0,50);
  fhPionRingPabs = 
    new TH1F("hPionRingPabs","Momentum associated with pion ring",100,0,50);
  fhMuonRingPabs = 
    new TH1F("hMuonRingPabs","Momentum associated with muon ring",100,0,50);

  fhElectronSpotPabs = 
    new TH1F("hElectronSpotPabs","Momentum associated with electron spot",100,0,50);
  fhPionSpotPabs = 
    new TH1F("hPionSpotPabs","Momentum associated with pion spot",100,0,50);
  fhMuonSpotPabs = 
    new TH1F("hMuonSpotPabs","Momentum associated with muon spot",100,0,50);


  fhElectronRingPabsPt = new TH2F("hfElectronRingPabsPt","Ring momentum vs pt",
				  50,0.,14.,50,0.,4.);
  fhElectronSpotPabsPt = new TH2F("hfElectronSpotPabsPt","MC ring momentum vs pt",
				  50,0.,14.,50,0.,4.);


  fhEff=new TH1F("hEff","Ring reconstruction efficiency",100,0,50);
  fhDist = new TH1F("hDist","Distance between ring and nearest LighSpot",100,-1.,5.);

  fhNonRecDist=new TH1F("hNonRecDist","Distance between the centers of non-rec. rings",100,0.,20.);
  fhRecDist=new TH1F("hRecDist","Distance between the centers of rec. rings",100,0.,20.);
  fhSpotDist=new TH1F("hSpotDist","Distance between the centers of MC rings",100,0.,20.);

  fhRefElectronRingPabs = 
    new TH1F("hRefElectronRingPabs","Momentum assigned to rec. e^{+-} ring from target region",100,0,50);
  fhRefElectronSpotPabs = 
    new TH1F("hRefElectronSpotPabs","Momentum assigned to MC e^{+-} ring from target region",100,0,50);

  fhRefPionRingPabs = 
    new TH1F("hRefPionRingPabs","Momentum assigned to rec. #pi^{+-} ring from target region",100,0,50);
  fhRefPionSpotPabs = 
    new TH1F("hRefPionSpotPabs","Momentum assigned to MC #pi{+-} ring from target region",100,0,50);

  fhNhitsElSpot = new TH1F("hNhitsElSpot","Number of hits per electron MC ring",
			   50,0.,100.);
  fhNhitsElRing = new TH1F("hNhitsElRing","Number of hits per rec. electron ring",
			   50,0.,100.);
  fhNhitsPiSpot = new TH1F("hNhitsPiSpot","Number of hits per #pi^{+-} MC ring",
			   50,0.,100.);
  fhNhitsPiRing = new TH1F("hNhitsPiRing","Number of hits per rec. #pi^{+-} ring",
			   50,0.,100.);


  fhRadElSpot = new TH1F("hRadElSpot","Radius of e^{+-} MC ring",
			   50,0.,10.);
  fhRadElRing = new TH1F("hRadElRing","Radis of reconstructed e^{+-} ring",
			   50,0.,10.);
  fhRadPiSpot = new TH1F("hRadPiSpot","Radius of #pi^{+-} MC ring",
			   50,0.,10.);
  fhRadPiRing = new TH1F("hRadPiRing","Radius of reconstructed #pi^{+-} ring",
			   50,0.,10.);


  fhElDeltaRad = new TH1F("hElDeltaRad","R_{rec}-R_{MC} for e^{+-}",100,-5.,5.);
  fhPiDeltaRad = new TH1F("hPiDeltaRad","R_{rec}-R_{MC} for pi^{+-}",100,-5.,5.);


  fhGhostXY = new TH2F("hGhostXY","Ghosts XY",500,-163,163,850,-255,255);
  fhRingXY = new TH2F("hRingXY","Found rings XY",500,-163,163,850,-255,255);

//   fhGhostXY = new TH2F("hGhostXY","Ghosts XY",125,-163,163,213,-255,255);
//   fhRingXY = new TH2F("hRingXY","Found rings XY",125,-163,163,213,-255,255);

}

//------------------------------------------------------------------------------
CbmRichAnalysisRingFinder::~CbmRichAnalysisRingFinder()
{
  //Write histograms
  fhRingPabs->Sumw2(); fhRingPabs->Write();

  fhElectronRingPabs->Sumw2(); fhElectronRingPabs->Write();
  fhPionRingPabs->Sumw2(); fhPionRingPabs->Write();
  fhMuonRingPabs->Sumw2(); fhMuonRingPabs->Write();

  fhSpotPabs->Sumw2(); fhSpotPabs->Write();
  fhElectronSpotPabs->Sumw2(); fhElectronSpotPabs->Write();
  fhPionSpotPabs->Sumw2(); fhPionSpotPabs->Write();
  fhMuonSpotPabs->Sumw2(); fhMuonSpotPabs->Write();

  fhElectronRingPabsPt->Sumw2(); 
  fhElectronRingPabsPt->Write();
  fhElectronSpotPabsPt->Sumw2(); 
  fhElectronSpotPabsPt->Write();

  fhEff->Divide(fhRingPabs,fhSpotPabs);
  fhEff->Write();

  fhDist->Write();
  fhNonRecDist->Write();
  fhRecDist->Write();
  fhSpotDist->Write();

  fhRefElectronRingPabs->Sumw2(); fhRefElectronRingPabs->Write();
  fhRefElectronSpotPabs->Sumw2(); fhRefElectronSpotPabs->Write();

  fhRefPionRingPabs->Sumw2(); fhRefPionRingPabs->Write();
  fhRefPionSpotPabs->Sumw2(); fhRefPionSpotPabs->Write();

  fhNhitsElSpot->Sumw2();
  fhNhitsElSpot->Scale(1./fhNhitsElSpot->Integral());
  fhNhitsElSpot->Write();

  fhNhitsElRing->Sumw2();
  fhNhitsElRing->Scale(1./fhNhitsElRing->Integral());
  fhNhitsElRing->Write();

  fhNhitsPiSpot->Sumw2();
  fhNhitsPiSpot->Scale(1./fhNhitsPiSpot->Integral());
  fhNhitsPiSpot->Write();

  fhNhitsPiRing->Sumw2();
  fhNhitsPiRing->Scale(1./fhNhitsPiRing->Integral());
  fhNhitsPiRing->Write();

  fhRadElSpot->Sumw2();
  fhRadElSpot->Scale(1./fhRadElSpot->Integral());
  fhRadElSpot->Write();

  fhRadElRing->Sumw2();
  fhRadElRing->Scale(1./fhRadElRing->Integral());
  fhRadElRing->Write();

  fhRadPiSpot->Sumw2();
  fhRadPiSpot->Scale(1./fhRadPiSpot->Integral());
  fhRadPiSpot->Write();

  fhRadPiRing->Sumw2();
  fhRadPiRing->Scale(1./fhRadPiRing->Integral());
  fhRadPiRing->Write();

  fhGhostXY->Sumw2(); fhGhostXY->Write();
  fhRingXY->Sumw2();  fhRingXY->Write();

  fhElDeltaRad->Write();
  fhPiDeltaRad->Write();

}

//------------------------------------------------------------------------------
InitStatus CbmRichAnalysisRingFinder::Init()
{
  CbmRootManager *manager= CbmRootManager::Instance();

  // all tracks
  fListStack      = (TClonesArray*)manager->ActivateBranch("MCTrack");

  //RICH hits
  fListRICHpoints   = (TClonesArray*)manager->ActivateBranch("RichPoint");
 
  //RICH hits
  fListRICHhits   = (TClonesArray*)manager->ActivateBranch("RichHit");
 
  //RICH rings
  fListRICHrings  = (TClonesArray*)manager->ActivateBranch("RichRing");

  return kSUCCESS;

}

//------------------------------------------------------------------------------
void CbmRichAnalysisRingFinder::MakeLightSpots(TClonesArray& spots)
{

  spots.Clear();

  CbmRichHit* hit=NULL;

  Int_t nLightSpots = 0;
  Int_t iLightSpot;

  for(Int_t iHit=0; iHit<fListRICHhits->GetEntries(); iHit++) {
    hit = (CbmRichHit*)fListRICHhits->At(iHit); // RICH hit
    if (hit->GetRefIndex() < 0) continue;  // fake hit
    CbmRichPoint* point = (CbmRichPoint*)fListRICHpoints->At(hit->GetRefIndex());
    CbmMCTrack* hitTrack = (CbmMCTrack*)fListStack->At(point->GetTrackID());
    if(hitTrack->GetPdgCode() != 50000050) continue; // select only Cherenkov photons

    // Fill array with cherenkov photon parents
    Int_t  motherTrack = hitTrack->GetMotherID();
    Bool_t dejaVue = kFALSE;
    CbmRichLightSpot *spot;
    for (iLightSpot=0; iLightSpot<nLightSpots; iLightSpot++) {
      spot = (CbmRichLightSpot*)spots.At(iLightSpot);
      if (spot->GetTrackID() == motherTrack) {
	dejaVue = kTRUE;
	spot->AddHitIndex(iHit);
	break;
      }
    }
    if (!dejaVue) {
      spot = new(spots[nLightSpots]) CbmRichLightSpot();
      spot->SetTrackID(motherTrack);
      spot->AddHitIndex(iHit);
      nLightSpots++;
    }
  }

  for (iLightSpot=0; iLightSpot<nLightSpots; iLightSpot++)
    ((CbmRichLightSpot*)spots.At(iLightSpot))->RingFit();

  cout<<endl
      <<"================ found "
      <<spots.GetEntries()
      <<" LightSpots. "<<endl;

}

//------------------------------------------------------------------------------
void CbmRichAnalysisRingFinder::Exec(Option_t* option)
{

  //Produce LighSpots
  TClonesArray spots("CbmRichLightSpot",500);
  MakeLightSpots(spots);

  // makes analysis
  Int_t iSpot,iRing,pdg_code;
  Double_t p_spot,pt_spot,vx,vy,vz;
  CbmRichLightSpot* spot=NULL;
  CbmRichLightSpot* ring;
  Int_t nRecRings=0; // total number of reconstructed (>=70% hits from one track) rings
  TVector3 vertexP;

  TArrayI srings(10);
  DecayRings(113,srings,&spots,fListStack);

  TArrayI rings(10);
  DecayRings(113,rings,fListRICHrings,fListStack);
  
  printf("\t\trho0 --> %d spots, %d rings.\n",srings.GetSize(),rings.GetSize());

  //LightSpots (MC rings)
  for(iSpot=0; iSpot<spots.GetEntriesFast(); iSpot++) {
    spot = (CbmRichLightSpot*)spots.At(iSpot);
    p_spot = (((CbmMCTrack*)fListStack->At(spot->GetTrackID()))->GetMomentum()).Mag();
    pt_spot = (((CbmMCTrack*)fListStack->At(spot->GetTrackID()))->GetMomentum()).Pt();
    fhSpotPabs->Fill(p_spot);

    pdg_code = ((CbmMCTrack*)fListStack->At(spot->GetTrackID()))->GetPdgCode();
    vertexP = ((CbmMCTrack*)fListStack->At(spot->GetTrackID()))->GetStartVertex();
    vx = vertexP.X();
    vy = vertexP.Y();
    vz = vertexP.Z();

    if(abs(pdg_code)==11) // electron spot
      { 
	fhNhitsElSpot->Fill(spot->GetNhits());
	if(spot->GetRadius())
	  fhRadElSpot->Fill(spot->GetRadius());
	fhElectronSpotPabsPt->Fill(p_spot,pt_spot);
	fhElectronSpotPabs->Fill(p_spot); 
	if (fabs(vx)<0.1 && fabs(vy)<0.1 && fabs(vz)<0.1) // from target region 
	  fhRefElectronSpotPabs->Fill(p_spot);	
      }
    if(abs(pdg_code)==211) {
      if(spot->GetRadius())
	fhRadPiSpot->Fill(spot->GetRadius());
      fhNhitsPiSpot->Fill(spot->GetNhits());
      fhPionSpotPabs->Fill(p_spot); // pion spot
      if (fabs(vx)<0.1 && fabs(vy)<0.1 && fabs(vz)<0.1) // from target region 
	fhRefPionSpotPabs->Fill(p_spot);
    }
    if(abs(pdg_code)==13) fhMuonSpotPabs->Fill(p_spot); // muon spot
  }

  //found rings
  for(iRing=0; iRing<fListRICHrings->GetEntriesFast(); iRing++) {
    ring = (CbmRichLightSpot*)fListRICHrings->At(iRing);
    ring->RingFit();
    ring->Print("");
    Int_t track_id = ring->GetTrackID();
    Int_t pdg;
    if(track_id>=-1)
      pdg = ((CbmMCTrack*)fListStack->At(track_id))->GetPdgCode();
    else
      pdg = -333;
    printf("TrackID is %d, PDGcode %d\n\n",track_id,pdg);
  
    Int_t nearest_spot = -111;
    Double_t min_dist = 1.e+10;
    for(iSpot=0; iSpot<spots.GetEntriesFast(); iSpot++) {
      spot = (CbmRichLightSpot*)spots.At(iSpot);
      Double_t dx = spot->GetCenterX() - ring->GetCenterX();
      Double_t dy = spot->GetCenterY() - ring->GetCenterY();
      Double_t dist = sqrt(dx*dx+dy*dy);
      if(dist<min_dist) {
	min_dist = dist;
	nearest_spot = iSpot;
      }
    }

    fhDist->Fill(min_dist);
    fhRingXY->Fill(ring->GetCenterX(), ring->GetCenterY());

    //if ring has TrackID, assign it a momentum of this track
    //ring has TrackID, if >= 70% of its hits originates from one track
    //in this case it called "reconstructed"
    if(track_id<-1) {
      fhGhostXY->Fill(ring->GetCenterX(), ring->GetCenterY());
      continue; 
    }

    nRecRings++;
    p_spot = (((CbmMCTrack*)fListStack->At(track_id))->GetMomentum()).Mag();
    pt_spot = (((CbmMCTrack*)fListStack->At(track_id))->GetMomentum()).Pt();
    fhRingPabs->Fill(p_spot);
    pdg_code = ((CbmMCTrack*)fListStack->At(track_id))->GetPdgCode();
    vertexP = ((CbmMCTrack*)fListStack->At(track_id))->GetStartVertex();
    vx = vertexP.X();
    vy = vertexP.Y();
    vz = vertexP.Z();

    //R_rec - R_MC
    for(Int_t imc=0; imc<spots.GetEntriesFast(); imc++) {
      spot = (CbmRichLightSpot*)spots.At(imc);
      if(spot->GetTrackID() == track_id && spot->GetRadius() && ring->GetRadius()) {
	if(abs(pdg_code)==11) fhElDeltaRad->Fill(spot->GetRadius()- ring->GetRadius());
	if(abs(pdg_code)==211) fhPiDeltaRad->Fill(spot->GetRadius()- ring->GetRadius());
      }
    }

    if(abs(pdg_code)==11) {
      if(ring->GetRadius())
	fhRadElRing->Fill(ring->GetRadius()); 
      fhNhitsElRing->Fill(ring->GetNhits()); 
      fhElectronRingPabs->Fill(p_spot); // rec. electron ring
      fhElectronRingPabsPt->Fill(p_spot,pt_spot); // P vs Pt
      if (fabs(vx)<0.1 && fabs(vy)<0.1 && fabs(vz)<0.1) // from target region
	fhRefElectronRingPabs->Fill(p_spot);
    }
    if(abs(pdg_code)==211) {
      if(ring->GetRadius())
	fhRadPiRing->Fill(ring->GetRadius()); 
      fhNhitsPiRing->Fill(ring->GetNhits()); 
      fhPionRingPabs->Fill(p_spot); // rec. pion ring
      if (fabs(vx)<0.1 && fabs(vy)<0.1 && fabs(vz)<0.1) // from target region
	fhRefPionRingPabs->Fill(p_spot);
    }
    if(abs(pdg_code)==13) fhMuonRingPabs->Fill(p_spot); // rec. muon ring

  }

  // distance between ring centers
  Int_t pdg1 = -333;
  Int_t pdg2 = -334;
  Double_t dx,dy,dr;
  Int_t i,j;
  Double_t min_dist_rec=1.e+10;
  Double_t min_dist_gst=1.e+10;
  for (i=0; i<fListRICHrings->GetEntriesFast(); i++) {
    CbmRichRing* r1 = (CbmRichRing*)fListRICHrings->At(i);

    for (j=0; j<fListRICHrings->GetEntriesFast(); j++) {
      CbmRichRing* r2 = (CbmRichRing*)fListRICHrings->At(j);
      if(!r2->IsEqual(r1)) {
	dx = r1->GetCenterX() - r2->GetCenterX();
	dy = r1->GetCenterY() - r2->GetCenterY();
	dr = sqrt(dx*dx+dy*dy);
//	if( (!r1->Reconstructed()) && (!r2->Reconstructed()) ) 
	  if(dr<min_dist_gst) min_dist_gst = dr;
//	if(r1->Reconstructed() && r2->Reconstructed())
	  if(dr<min_dist_rec) min_dist_rec = dr;	  
      }
    }
    fhNonRecDist->Fill(min_dist_gst);
    fhRecDist->Fill(min_dist_rec);
  } 
   
  //distance between light spot centers
  Float_t min_dist_e = 10.; // e-rings are "well-separated", if dist>12cm
  TObjArray well_sep_e_spots(100); 
  TObjArray well_sep_pi_spots(100); 

  Double_t min_dist = 1.e+10;
  for (i=0; i<spots.GetEntriesFast(); i++) {
    CbmRichLightSpot* s1 = (CbmRichLightSpot*)spots.At(i);
    pdg1 = ((CbmMCTrack*)fListStack->At(s1->GetTrackID()))->GetPdgCode();
    for (j=0; j<spots.GetEntriesFast(); j++) {
      CbmRichLightSpot* s2 = (CbmRichLightSpot*)spots.At(j);
      pdg2 = ((CbmMCTrack*)fListStack->At(s2->GetTrackID()))->GetPdgCode();
      dx = s1->GetCenterX() - s2->GetCenterX();
      dy = s1->GetCenterY() - s2->GetCenterY();
      dr = sqrt(dx*dx+dy*dy);
      if(dr==0.) continue;
      if(dr<min_dist) min_dist=dr;  
    }
    fhSpotDist->Fill(min_dist);

    if(min_dist>min_dist_e) {
      if(abs(pdg1)==11)
	well_sep_e_spots.Add(s1);
      if(abs(pdg1)==211)
	well_sep_pi_spots.Add(s1);
    }

  }

  printf("\n\t\tReconstructed rings: %d, total: %d, found: %d\n",
	 nRecRings,spots.GetEntriesFast(),fListRICHrings->GetEntriesFast());
  printf("\t\t%d well separated e+- spots, %d pi+-\n",
	 well_sep_e_spots.GetEntriesFast(),
	 well_sep_pi_spots.GetEntriesFast());
    
}

//------------------------------------------------------------------------------
void CbmRichAnalysisRingFinder::Finish()
{}

//------------------------------------------------------------------------------
void CbmRichAnalysisRingFinder::DecayRings(Int_t pdgCode,
					   TArrayI& rings, 
					   TClonesArray* ringpool,
					   TClonesArray* trackpool)
{
  //Found rings from primary vector meson (i.e rho0, pdgCode=113) decay electrons
  //and place indices (in ringpool array) into &rings.

  CbmMCTrack* trk;
  Int_t pdg, mid, mpdg, tr;
  Int_t nFound=0;

  for(Int_t rn=0; rn<ringpool->GetEntriesFast(); rn++) {
    tr = ((CbmRichLightSpot*)ringpool->At(rn))->GetTrackID();
    if(tr<0) continue;

    trk = (CbmMCTrack*)trackpool->At(tr);
    pdg=trk->GetPdgCode();
    mid=trk->GetMotherID();

    if(mid==-1) {
      if(abs(pdg)==11) mpdg=pdgCode;
    }
    else
      mpdg=((CbmMCTrack*)trackpool->At(mid))->GetPdgCode();

    if(abs(pdg)==11&&mpdg==pdgCode) {
      rings.AddAt(rn,nFound);
      nFound++;
    }
  }
  rings.Set(nFound);
}


//------------------------------------------------------------------------------
ClassImp(CbmRichAnalysisRingFinder)