//-----------------------------------------------------------
// 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"
#include <cmath>
// C/C++ Headers ----------------------
#include <map>
#include <set>
#include <algorithm>
#include <string>
#include <assert.h>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "TrackFitStat.h"
#include "GFTrackCand.h"
#include "GFAbsRecoHit.h"
#include "PndTpcPlanarRecoHit.h"
#include "PndTpcPoint.h"
#include "PndMCTrack.h"
#include "TMath.h"
#include "TH1D.h"
#include "TDatabasePDG.h"
#include "TVector2.h"
#include "MCTruthAnnex.h"
#include "LSLTrackRep.h"
#include "GeaneTrackRep.h"
#include "RKTrackRep.h"
#include "PndTpcSPHit.h"

using std::cout;
using std::cerr;
using std::endl;

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


TrackFitStatTask::~TrackFitStatTask()
{
}


InitStatus
TrackFitStatTask::Init()
{
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::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<unsigned int,int> mctruthmap; // index in array mapped to 0,1,2
                                         // 0 -> not valid
                                         // 1 -> valid but not found
                                         // 2 -> valid, found
  unsigned int nmc=_mcTrackArray->GetEntriesFast();
  unsigned int ntpcMChits=_mcPndTpcHitArray->GetEntriesFast();
  
  //nmctrks=1; 	///////////////////////!GetFirstTrackOnly!!!!!!!!
   for(unsigned int imc=0; imc<nmc;++imc)
  { // loop over mc-tracks
  	 bool bValidTrack=true; 
	// check if this mctrack has created enough hits in tpc
	int counter=0;
	double length=0;
	TVector2 oldpoint;
	for(unsigned 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;
	PndMCTrack* mc=(PndMCTrack*)_mcTrackArray->At(imc);
	
	// create MCTruth Annex
	MCTruthAnnex* annex=new((*_mcAnnexArray)[imc]) MCTruthAnnex();
	annex->setPndTpcLength(length);
	annex->setPndTpcHits(counter);
	
	
	if(counter<_minPndTpcHits)	{
		bValidTrack=false;
	}
	_trklengthtpc->Fill(length);
	// check if this track is inside the valid sample
	double mcq=0;
	if(TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode()))	{
			mcq=TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode())->Charge()/3.;
	}
	else	{
			cout << "TrackFitStatTask: Can' t get Particle for PDG " << mc->GetPdgCode() << endl;
		}
	
	if(mcq==0)	{
	  //cerr<<"Not charged"<<endl;
	    bValidTrack=false;
	}	 	
	if(_pdgselect && mc->GetPdgCode()!=_pdgId)	{
	  //cerr<<"Not selected PDGId"<<endl;
		bValidTrack=false;
	}
	double pmc=mc->Get4Momentum().P();				
	if(pmc<_pmin || _pmax<pmc) {
	  //cerr<<"Not in selected momentum range"<<endl;
		bValidTrack=false;										
	}
	double theta_mc=mc->Get4Momentum().Theta();		
	if(theta_mc<_thetamin || _thetamax<theta_mc)	{
	  //cerr<<"Not in selected angular range"<<endl;
	 	bValidTrack=false;
	 }
	
	// if we are here the track lies inside our valid sample
	// -> put it into the map
	if(bValidTrack)	{	
		mctruthmap[imc]=1;	//i don' t think it will change anything, to put a 1 here
		std::cout<<"Valid MCTrack found with p="<<pmc<<" GeV/c"
		<<"  PDGId="<<mc->GetPdgCode()
		<<"  q="<<mcq
		<<"  nPndTpcHits="<<annex->getPndTpcHits()
		<<"  theta="<<theta_mc<<std::endl;
	}
	else	{
		mctruthmap[imc]=0;
	}
  }// end loop over mc tracks



  Int_t nTracks=_trackArray->GetEntriesFast();
  
  for(Int_t i=0; i<nTracks; ++i)
  { // loop over tracks
	GFTrack* track=(GFTrack*)_trackArray->At(i);
	//clumsily extract failed hits of cardinal rep and pass to TrackFitStat
	const GFTrackCand& cand = track->getCand();

	std::set<unsigned int> detIDs = cand.GetUniqueDetIDs();
	std::set<unsigned int>::iterator iter;
	int NFH=0;

	for(iter=detIDs.begin(); iter!=detIDs.end(); ++iter) {
	  int detID = *iter;

	  NFH+=track->getFailedHits(0);

	}	

	std::cout<<"TrackFitStatTask: recotrack has "<<NFH<<" failed hits"<<std::endl;

	if(track->getTrackRep(0)->getStatusFlag()!=0){
		std::cout<<"Trackfit not successful!"<<std::endl;
		continue;
	}
	Int_t size=_fitstatArray->GetEntriesFast();
	TrackFitStat* stat=new ((*_fitstatArray)[size]) TrackFitStat();
	
	//decide in which trackrep we are in
	GFAbsTrackRep* trackRep = track->getTrackRep(0);
	
	//LSL rep
	bool LSLREP = dynamic_cast<LSLTrackRep*>(trackRep);
	bool GEANEREP = dynamic_cast<GeaneTrackRep*>(trackRep);
	bool RKREP = dynamic_cast<RKTrackRep*>(trackRep);

	if(_doRes){
	  // fill tpc residuals
	  std::vector<double> res;
	  std::vector<double> sigmas;
	  track->getResiduals(2,0,0,res);//,sigmas);
	  stat->fillPndTpcResX(res);
	  //stat->fillPndTpcSigX(sigmas);
	  res.clear();
	  sigmas.clear();
	  track->getResiduals(2,1,0,res);//,sigmas);
	  stat->fillPndTpcResY(res);
	  //stat->fillPndTpcSigY(sigmas);
	  
	  
	  std::vector<unsigned int> clustersize;
	  std::vector<double> clusteramp;
	  unsigned int nh=track->getNumHits();

	  for(unsigned int ih=0; ih<nh; ++ih){ //loop over hits
	    GFAbsRecoHit* abshit=track->getHit(ih);
	    if(LSLREP) {
	      PndTpcPlanarRecoHit* tpchit=dynamic_cast<PndTpcPlanarRecoHit*>(abshit);
	       if(tpchit!=NULL){
		 clustersize.push_back(tpchit->cluster_size());
		 clusteramp.push_back(tpchit->cluster_amp());
	       }
	    }
	    // if(GEANEREP) {
// 	      PndTpcSPHit* tpchit=dynamic_cast<PndTpcSPHit*>(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, palt, pstart, q;
	if(LSLREP){
	  p=track->getMom().Mag();
	  //palt=1./fabs(track->getTrackRep(0)->getState()[4][0]);
	  //pstart=1.0; 
	  q=track->getCharge();
	  pstart=0;
	  std::cout<<"p="<<p<<"  q="<<q<<"  pstart="<<pstart<<std::endl;
	}
	if(GEANEREP){
	  p=track->getMom().Mag();
	  //palt=1./fabs(track->getTrackRep(0)->getState()[4][0]);
	  //pstart=1./fabs(track->getTrackRep(0)->getStartState()[4][0]);
	  pstart=0;
	  q=track->getCharge();
	
	  std::cout<<"p="<<p<<"  q="<<q<<"  pstart="<<pstart<<std::endl;
	}
	if(RKREP){
	  p=track->getMom().Mag();
	  //palt=1./fabs(track->getTrackRep(0)->getState()[4][0]);
	  //pstart=1./fabs(track->getTrackRep(0)->getStartState()[4][0]);
	  pstart=0;
	  q=track->getCharge();
	
	  std::cout<<"p="<<p<<"  q="<<q<<"  pstart="<<pstart<<std::endl;
	}
 


	//std::cout<<"\n *** TrackFitStatTask: Writing fit results to stat object ***\n"<<std::endl;
	stat->setp(p);
	stat->setmom(track->getMom());
	stat->setpstart(pstart);
	stat->setcharge(q);
	stat->addFailedHits(NFH);
	
	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
       	int index=0;
	for(unsigned int imc=0;imc<nmc;++imc){
	  //cerr << mctruthmap[imc] << endl;
	  PndMCTrack* mc=(PndMCTrack*)_mcTrackArray->At(imc);
	  double mcp=mc->Get4Momentum().P();
	  double mcq=-100;
	  if(TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode()))	{
			mcq=TDatabasePDG::Instance()->GetParticle(mc->GetPdgCode())->Charge()/3.;
		}
		else	{
			cout << "TrackFitStatTask: Can' t get Particle for PDG (2) " << mc->GetPdgCode() << endl;
		}
		//if(fabs(p-mcp)/mcp<_mcPCut*0.023*mcp){
#if 1
// 		if(fabs(p-mcp)<_precotol && !(mcit->second) )	{	//just see all misfit canditates
// 			cout << "TrackFitStatTask: PDG: " << mc->GetPdgCode() << "with q=" <<q
// 			<< " and mcq= " << mcq
// 			<< " and mcp= " <<  mcp 
// 			<< " passed Momentum cut..."
// 			<< " but actually is no valid track at all!" << endl;
// 			stat->setpmc(-10.0);
// 			stat->setp(-100.0);
// 		}

	//	else if(fabs(p-mcp)<0.01 && (mcit->second) && q==mcq )	{
	  if(fabs(p-mcp)<_precotol && (mctruthmap[imc]>0) )	{	
	//	if(mcit->second) 	{	//set properties of first track, momentum reconstruction not possible now
	    cout << "Setting mcp="<<mcp<<  endl;	//if it occurs twice the charge cut fails, too
	    stat->setpmc(mcp);
	    stat->setp(p);				
	    stat->setmccharge(mcq);
	    stat->setpdg(mc->GetPdgCode());
	    stat->setmotherid(mc->GetMotherID());
	    //assert(mc->GetMotherID()==-1); 
	    MCTruthAnnex* annex=(MCTruthAnnex*)_mcAnnexArray->At(index);
	    annex->setReco();
	    ++mctruthmap[imc];
	    //break;
	  }
	  else	if(mctruthmap[imc]>0)  {
	    stat->setpmc(mcp);	
	    stat->setp(-100.0);
	    stat->setmotherid(mc->GetMotherID());
	    stat->setpdg(mc->GetPdgCode());
	    stat->setmccharge(mcq);
	  }
	  else	{
	    stat->setpmc(-100.0);
	    stat->setp(-100.0);
	  }
#endif
	  ++index;
	} // end loop over mc-tracks
	
  } // end loop over tracks
  
  // analyse mc-truth association
  
  int foundcounter=0;
  int validcounter=0;
  for(unsigned int imc=0;imc<nmc;++imc)
    {
      //cerr << mctruthmap[imc] << endl;
      if(mctruthmap[imc]>0)++validcounter;	
      if(mctruthmap[imc]>1)++foundcounter;
    }
  
  std::cout<<"TrackFitStatTask:: "<<foundcounter<<" tracks of "<<validcounter<<" found"<<std::endl;

  if(validcounter>0){
    double eff=((double)foundcounter)/((double)(validcounter));
    _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)