//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TrackFitStatTask
//      see TrackFitStatTask.h for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sebastian Neubert    TUM            (original author)
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "TrackFitStatTask.h"

// C/C++ Headers ----------------------
#include <map>
#include <algorithm>
#include <string>
#include <assert.h>

// Collaborating Class Headers --------
#include "CbmRootManager.h"
#include "TClonesArray.h"
#include "Track.h"
#include "TrackFitStat.h"
#include "PndTpcPlanarRecoHit.h"
#include "PndTpcPoint.h"
#include "CbmMCTrack.h"
#include "TMath.h"
#include "TH1D.h"
#include "TDatabasePDG.h"
#include "TVector2.h"
#include "MCTruthAnnex.h"


TrackFitStatTask::TrackFitStatTask()
  : CbmTask("TrackFitStatistics"), _persistence(kFALSE), _doRes(kFALSE),
    _mcPCut(0.01), _pmin(0), _pmax(100), _thetamin(0), 
    _thetamax(TMath::TwoPi()), _minPndTpcHits(0), _pdgselect(false)
{
  _trackBranchName = "TrackPreFit";
  _mcBranchName = "MCTrack";
  _mcPndTpcBranchName = "PndTpcPoint";
}


TrackFitStatTask::~TrackFitStatTask()
{
}


InitStatus
TrackFitStatTask::Init()
{
  //Get ROOT Manager
  CbmRootManager* ioman= CbmRootManager::Instance();

  if(ioman==0)
    {
      Error("TrackReadTask::Init","RootManager not instantiated!");
      return kERROR;
    }
  
  // Get input collection
  _trackArray=(TClonesArray*) ioman->GetObject(_trackBranchName);
  
  if(_trackArray==0)
    {
      Error("TrackFitStatTask::Init","Track-array not found!");
      return kERROR;
    }


  _mcTrackArray=(TClonesArray*) ioman->GetObject(_mcBranchName);
  if(_mcTrackArray==0)
    {
      Error("TrackFitStatTask::Init","MC-Truth not found!");
      return kERROR;
    }

  _mcPndTpcHitArray=(TClonesArray*) ioman->GetObject(_mcPndTpcBranchName);
  if(_mcPndTpcHitArray==0)
    {
      Error("TrackFitStatTask::Init","PndTpc MC-Hits not found!");
      return kERROR;
    }


  // Create output collection
  _fitstatArray = new TClonesArray("TrackFitStat");
  ioman->Register("TrackFitStat","GenFit",_fitstatArray,_persistence);

  _mcAnnexArray = new TClonesArray("MCTruthAnnex");
  ioman->Register("MCTruthAnnex","GenFit",_mcAnnexArray,_persistence);

  _efficiency=new TH1D("efficiency","tracking efficiency",110,0,1.1);
  _trklengthtpc=new TH1D("tracklengthtpc","track length in tpc",150,0,150);

  return kSUCCESS;

}

 void 
TrackFitStatTask::SetMCCuts(Double_t pmin, Double_t pmax, 
			    Double_t thetamin, Double_t thetamax,
			    Int_t minPndTpcHits)
{
  _pmin=pmin;_pmax=pmax;
  _thetamin=thetamin;_thetamax=thetamax;
  _minPndTpcHits=minPndTpcHits;
}


void
TrackFitStatTask::Exec(Option_t* opt)
{
  std::cout << "TrackFitStatTask::Exec" << std::endl;

  // clear output
  if(_fitstatArray==0) Fatal("TrackFitStat::Exec)","No FitStat OutputArray");
  _fitstatArray->Delete();
  if(_mcAnnexArray==0) Fatal("TrackFitStat::Exec)","No McAnnex OutputArray");
  _mcAnnexArray->Delete();

  // prepare MonteCarlo Truth info
  std::map<CbmMCTrack*,int> mctruthmap;
  int nmctrks=_mcTrackArray->GetEntriesFast();
  int ntpcMChits=_mcPndTpcHitArray->GetEntriesFast();
  for(int imc=0; imc<nmctrks;++imc){ // loop over mc-tracks
    // check if this mctrack has created enough hits in tpc
    int counter=0;
    double length=0;
    TVector2 oldpoint;
    for(int ih=0; ih<ntpcMChits;++ih){
      PndTpcPoint* point=(PndTpcPoint*)_mcPndTpcHitArray->At(ih);
      if(point->GetTrackID()==imc){
	++counter;
	if(ih>0){
	  TVector2 point2(point->GetX(),point->GetY());
	  length+=(point2-oldpoint).Mod();
	}
	oldpoint.Set(point->GetX(),point->GetY());
      }
      
      //if(counter>=_minPndTpcHits)break;
    }
    //std::cout<<"MCTrack"<<imc<<" has "<<counter<<" hits in TPC."<<std::endl;
    CbmMCTrack* mc=(CbmMCTrack*)_mcTrackArray->At(imc);

    // create MCTruth Annex
    MCTruthAnnex* annex=new((*_mcAnnexArray)[imc]) MCTruthAnnex();
    annex->setPndTpcLength(length);
    annex->setPndTpcHits(counter);
    

    
    if(counter<_minPndTpcHits)continue;
    _trklengthtpc->Fill(length);
    // check if this track is inside the valid sample
    double mcq=TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode())->Charge()/3.;
    if(mcq==0)continue;
    if(_pdgselect && mc->GetPdgCode()!=_pdgId)continue;
    double pmc=mc->Get4Momentum().P();
    if(pmc<_pmin || _pmax<pmc) continue;
    double theta_mc=mc->Get4Momentum().Theta();
    if(theta_mc<_thetamin || _thetamax<theta_mc) continue;
    std::cout<<"Valid MCTrack found with p="<<pmc<<" GeV/c"
	     <<"  PDGId="<<mc->GetPdgCode()
	     <<"  q="<<mcq
	     <<"  nPndTpcHits="<<annex->getPndTpcHits()
	     <<"  theta="<<theta_mc<<std::endl;

    // if we are here the track lies inside our valid sample
    // -> put it into the map
    mctruthmap[mc]=0;
  }



  Int_t nTracks=_trackArray->GetEntriesFast();
  for(Int_t i=0; i<nTracks; ++i){ // loop over tracks
    Track* track=(Track*)_trackArray->At(i);


    Int_t size=_fitstatArray->GetEntriesFast();
    TrackFitStat* stat=new ((*_fitstatArray)[size]) TrackFitStat();

    if(track->getTrackRep(0)->getStatusFlag()!=0){
      std::cout<<"Trackfit not successfull!"<<std::endl;
      continue;
    }

    if(_doRes){
      // fill tpc residuals
      std::vector<double> res;
    
      track->getResiduals(2,0,0,res);
      stat->fillPndTpcResX(res);
      res.clear();
      track->getResiduals(2,1,0,res);
      stat->fillPndTpcResY(res);
    

    std::vector<unsigned int> clustersize;
    std::vector<unsigned int> clusteramp;
    unsigned int nh=track->getNumHits();
    for(unsigned int ih=0; ih<nh; ++ih){ //loop over hits
      AbsRecoHit* abshit=track->getHit(ih);
      PndTpcPlanarRecoHit* tpchit=dynamic_cast<PndTpcPlanarRecoHit*>(abshit);
      if(tpchit!=NULL){
	clustersize.push_back(tpchit->cluster_size());
	clusteramp.push_back(tpchit->cluster_amp());
      }
    }
    if(clustersize.size()==0)std::cout<<"No TPCHits found!"<<std::endl;
    stat->fillPndTpcClusterSize(clustersize);
    stat->fillPndTpcClusterAmp(clusteramp);
    }    


    double p=track->getMom().Mag();
    double palt=1./fabs(track->getTrackRep(0)->getState()[4][0]);
    double pstart=1./fabs(track->getTrackRep(0)->getStartState()[4][0]);
    double q=track->getCharge();

    std::cout<<"p="<<p<<"  q="<<q<<"  pstart="<<pstart<<std::endl;


    stat->setp(p);
    stat->setmom(track->getMom());
    stat->setpstart(pstart);
    stat->setcharge(q);

    double s=track->getTrackRep(0)->getCov()[4][4];
    if(s>0){
      double sigp=p*p*sqrt(s);
      stat->setsigp(sigp);
    }

    // try to associate a MonteCarlo Track
    std::map<CbmMCTrack*,int>::iterator mcit=mctruthmap.begin();
    int index=0;
    while(mcit!=mctruthmap.end())
      {
      CbmMCTrack* mc=mcit->first;
      double mcp=mc->Get4Momentum().P();
      double mcq=TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode())->Charge()/3.;
      if(fabs(p-mcp)/mcp<_mcPCut*0.023*mcp){
	stat->setpmc(mcp);
	stat->setmccharge(mcq);
	MCTruthAnnex* annex=(MCTruthAnnex*)_mcAnnexArray->At(index);
	annex->setReco();
	++(mcit->second);
	break;
      }
      ++mcit;
      ++index;
    } // end loop over mc-tracks

  } // end loop over tracks
  
  // analyse mc-truth association
  std::map<CbmMCTrack*,int>::iterator mcit=mctruthmap.begin();
  int counter=0;
  while(mcit!=mctruthmap.end())
    {
      if(mcit->second>0)++counter;
      ++mcit;
    }

  std::cout<<"TrackFitStatTask:: "<<counter<<" tracks of "<<mctruthmap.size()<<" found"<<std::endl;

  if(mctruthmap.size()>0){
    double eff=((double)counter)/((double)(mctruthmap.size()));
    _efficiency->Fill(eff);
  }
  

  return;
}

void 
TrackFitStatTask::WriteHistograms(const TString& filename){
  TFile* file = new TFile(filename,"UPDATE");
  file->mkdir("TrackFitStat");
  file->cd("TrackFitStat");

  _efficiency->Write();
  delete _efficiency;
  _efficiency=NULL;

  _trklengthtpc->Write();
  delete _trklengthtpc;
  _trklengthtpc=NULL;

  file->Close();
  delete file;
}


ClassImp(TrackFitStatTask)