//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcdEdxTask
//      see TpcdEdxTask.hh for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer, Christian Hoeppner     TUM         
//
//
//-----------------------------------------------------------

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

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

// C/C++ Headers ----------------------
#include <cmath>
#include <iostream>
#include <map>
#include <assert.h> 
#include <stdexcept>

#include "../base/TpcdEdx.h"
// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "TpcCluster.h"
#include "TpcDigiMapper.h"
#include "TpcFrontend.h"
#include "TpcSPHit.h"
#include "PndMCTrack.h"
#include "TpcPoint.h"
#include "TpcAlignmentManager.h"

#include "TpcPadShapePool.h"
#include "TpcPadPlane.h"
#include "TpcAbsPadShape.h"
#include "TpcPad.h"

#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "TpcDigiPar.h"


#include "GFAbsTrackRep.h"
#include "GFRecoHitFactory.h"
#include "GFKalman.h"
#include "GFException.h"
#include "TH1D.h"
#include "TFile.h"
#include "TGeoTrack.h"

#include "GFDetPlane.h"
#include "RecoHits/GFAbsRecoHit.h"
#include "TVector3.h"

//#include <signal.h>
//#include <stdlib.h>

#include "PndDetectorList.h"

// Class Member definitions -----------


TpcdEdxTask::TpcdEdxTask()
: FairTask("dE/dx Task"), _persistence(kFALSE),
  _DX(1.), _idealdEdx(kFALSE), fUseFirstDigiPos(kFALSE),
  fIgnoreEdgeDigis(kFALSE)
{
	_trackBranchName = "TpcTrackPostFit";
	_clusterBranchName = "TpcCluster";
	_digiBranchName = "TpcDigi";
	_outBranchName = "TpcdEdx";
	//_spHitBranchName = "TpcSPHit";
	fVerbose = 0;
	//fPadPlaneFile = "tpc/FOPI/par/padplane_FOPI.dat";
	//fShapeFile = "tpc/parfiles/PadShapePrototype.dat";
}


TpcdEdxTask::~TpcdEdxTask()
{
}

InitStatus
TpcdEdxTask::Init()
{
	//Get ROOT Manager
	FairRootManager* ioman= FairRootManager::Instance();

	if(ioman==0)
	{
		Error("TpcdEdxTask::Init","RootManager not instantiated!");
		return kERROR;
	}

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

	if(_idealdEdx) {
		_mcTrackArray=(TClonesArray*) ioman->GetObject("MCTrack");
		if(_mcTrackArray==0)
			Fatal("TpcdEdxTask::Init","MCTrack-array not found! Cannot calculate MC dEdx");

		_pointArray=(TClonesArray*) ioman->GetObject("TpcPoint");
		if(_pointArray==0)
			Fatal("TpcdEdxTask::Init","Point-array not found! Cannot calculate MC dEdx");
	}

	_clusterArray=(TClonesArray*) ioman->GetObject(_clusterBranchName);
	if(_clusterArray==0)
	{
		Error("TpcdEdxTask::Init","Cluster/Hit - array not found!");
		return kERROR;
	}

	_digiArray=(TClonesArray*) ioman->GetObject(_digiBranchName);
	if(_digiArray==0)
	{
		Error("TpcdEdxTask::Init","Digi-array not found!");
		return kERROR;
	}

	_dEdxOutArray = new TClonesArray("TpcdEdx");
	ioman->Register(_outBranchName,"TpcdEdx",_dEdxOutArray,_persistence);

	_dEdxMCOutArray = new TClonesArray("TpcdEdx");
	ioman->Register("dEdx_MC","Tpc",_dEdxMCOutArray,_persistence);

	// setup histograms
	_distHist=new TH1D("distHist","Stepping lentghs",500,-3.,3.);
	_dirHist=new TH1D("dirHist","direction abundances",10,-2.,2.);

	_hitBranchID = FairRootManager::Instance()->GetBranchId(_clusterBranchName);

	std::cout<<"Initializing PadPlane ... "<<"\n";;
	//TpcPadShapePool* pool = fPar->getPadShapes();
	fPlane = fPar->getPadPlane();
	fRMin = fPar->getRMin();
	fRMax = fPar->getRMax();

	if(fIgnoreEdgeDigis) {
		fRMin = fRCutoff1;
		fRMax = fRCutoff2;
	}

	std::cout<<"TpcdEdxTask::Init: success\n";
	return kSUCCESS;
}

void
TpcdEdxTask::SetParContainers() {

	std::cout<<"TpcClusterFinderTask::SetParContainers"<<"\n";
	std::cout.flush();

	// Get run and runtime database
	FairRun* run = FairRun::Instance();
	if ( ! run ) Fatal("SetParContainers", "No analysis run");

	FairRuntimeDb* db = run->GetRuntimeDb();
	if ( ! db ) Fatal("SetParContainers", "No runtime database");

	// Get Tpc digitisation parameter container
	fPar= (TpcDigiPar*) db->getContainer("TpcDigiPar");
	if (! fPar ) Fatal("SetParContainers", "TpcDigiPar not found");

}

void
TpcdEdxTask::Exec(Option_t* opt)
{
	if (fVerbose) std::cout<<"TpcdEdxTask::Exec()"<<"\n";
	FairRootManager* ioman= FairRootManager::Instance();

	static unsigned int nTotalDigis(0);
	static unsigned int nDoubleDigis(0);
	int evts(0);

	_dEdxOutArray->Delete();
	_dEdxMCOutArray->Delete();


	Int_t ntracks=_trackArray->GetEntriesFast();

	if(ntracks>20000){
		std::cout<<"ntracks="<<ntracks<<" Evil Event! skipping"<<"\n";
		return;
	}

	if (fVerbose) std::cout<<"start loop over "<<ntracks<<" tracks\n";
	for(unsigned int itr=0;itr<ntracks;++itr){
		double totDist = 0.;
		if (fVerbose) std::cout<<"TpcdEdxTask::Exec(): starting track "<<itr<<"\n";
		GFTrack* trk=(GFTrack*)_trackArray->At(itr);
		if (fVerbose) std::cout<<"*** Number of hits in track: "<<trk->getNumHits()<<" ***"<<"\n";

		TpcdEdx dedx;
		TpcdEdx dedx_MC;

		GFAbsTrackRep* theRep = trk->getCardinalRep();
		if (fVerbose) std::cout<<"check status flag of track "<< theRep->getStatusFlag()<<"\n";
		if (theRep->getStatusFlag() != 0) {
			// write empty dEdx and continue with next trk
			new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

			if(_idealdEdx) {
				new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
			}
			continue;
		}


		GFTrackCand cand = trk->getCand();

		bool useSPHits = false;
		if(TString(_clusterArray->GetClass()->GetName()) == "TpcSPHit")
			useSPHits = true;

		//these are SORTED, first entry corresponds to track's start-hit
		std::vector<int> hitIDs = cand.getHitIDs(_hitBranchID);

		//GET MC INFORMATION FOR CROSS-CHECK
		//only works with IdealTracking and pure trackIDs
		std::vector<TpcPoint*> pointlist;
		if(_idealdEdx) {
			if (fVerbose) std::cerr<<"collect MC points"<<"\n";
			//loop over MC points and collect
			for(int p=0; p<_pointArray->GetEntriesFast(); ++p) {
				int id = ((TpcPoint*)_pointArray->At(p))->GetTrackID();
				if(id==0) //only primary tracks
					pointlist.push_back((TpcPoint*)_pointArray->At(p));
			}
		}

		TVector3 pos,mom;
		TVectorD protState = theRep->getState();
		if (fVerbose) std::cerr<<"get pos and mom"<<"\n";
		try{
			pos = theRep->getPos();
			mom = theRep->getMom();
		}
		catch(GFException& e){
			e.what();
			// write empty dEdx and continue with next trk
			new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);
			if(_idealdEdx) {
				new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
			}
			continue;
			//TODO: exception handling
		}

		TVector3 startDir = mom;
		startDir.SetMag(1.);

		TVector3 altFirstPos;

		if (fVerbose) std::cout<<"get reference plane\n";
		//reference plane:
		GFDetPlane pl;
		pl.setO(pos);
		pl.setNormal(startDir);

		if (fVerbose) std::cerr<<"collect cluster"<<"\n";
		// get all digis in the track
		std::map<unsigned int, TpcDigi*> digis;   //<TCA index, pointer>
		unsigned int count = 0;
		TClonesArray* actClArr;
		int foundDigi = -1; //first digi in (opposite) track direction

		if (fVerbose) std::cout<<"start loop over "<<hitIDs.size()<<" clusters\n";
		for(int i=0; i<hitIDs.size(); i++){ //loop over hits (clusters)
			TpcCluster* cl;
			if(useSPHits) {
				if (fVerbose) std::cout<<"using sphits\n";
				TpcSPHit* theHit = (TpcSPHit*) _clusterArray->At(hitIDs[i]);
				int clBrID = theHit->getClusterBranchID();
				int clID = theHit->getClusterID();
				if(count==0) {
					TString clBrName = ioman->GetBranchName(clBrID);
					actClArr = (TClonesArray*) ioman->GetObject(clBrName);
					if (fVerbose) std::cout<<"ID:"<<clBrID<<" clbrName: "<<clBrName<<"\n";
				}
				if (fVerbose) std::cout<<"get cluster "<<clID<<"\n";
				cl = (TpcCluster*)(actClArr->At(clID));
				count++;
			}
			else{
				if (fVerbose) std::cout<<"get cluster\n";
				cl = (TpcCluster*) _clusterArray->At(hitIDs[i]);
			}
			if (fVerbose) std::cout<<"get digimap\n";
			const std::map<unsigned int,double>* digiMap = cl->getDigiMap();

			double maxDist = 0.;  //(cm)

			if (fVerbose) std::cout<<"loop over "<<digiMap->size()<<" digis \n";
			for(std::map<unsigned int, double>::const_iterator it=digiMap->begin();it!=digiMap->end();++it){
				//avoid double-counting
				if(!digis.count(it->first))
					digis[it->first] = (TpcDigi*) _digiArray->At(it->first);
				else
					continue;
				TpcDigi* aDigi = digis[it->first];

				//only for first hit (cluster): find first Digi
				if(i==0 && fUseFirstDigiPos) {
					//get Digi coordinates:
					TVector3 dPos;
					try {
						TpcDigiMapper::getInstance()->map(aDigi,dPos);
					}
					catch (...) {
						std::cerr<<"TpcdEdxTask::Could not get Position of digi! Skipping ..."<<"\n";
						continue;
					}
					dPos = TpcAlignmentManager::getInstance()->localToMaster("tpc",dPos);
					TVector3 pc, dpc;
					try{
						theRep->extrapolateToPoint(dPos,pc,dpc);
					}catch(GFException exp){
						std::cerr<<"TpcdEdxTask::Exec() error during first digi extrapolation:"<<"\n";
						std::cerr<<exp.what()<<"\n";
						continue;
					}
					TVector3 dDist = pl.getO() - pc;
					if(dDist.Mag() > maxDist && (dDist*pl.getNormal()) < 0.) {  //we want negative extrap!
						if(fIgnoreEdgeDigis)
							if(dPos.Perp() < fRCutoff1 || dPos.Perp() > fRCutoff2)
								continue;
						foundDigi = (*it).first;
						altFirstPos = dPos;
						maxDist = dDist.Mag();
					}
				}
			} //end loop over this cluster's digis
			if (fVerbose) std::cout<<"done looping over digis\n";

		} //end loop over hits (clusters)
		if (fVerbose) std::cout<<"done looping over clusters\n";

		//startDir *= -1.;
		dedx.setMom(mom);

		//if (fVerbose) {
		//	TVector3 pos0(((TpcCluster*)(_clusterArray->At(hitIDs[0])))->pos());
		//	std::cout<<"pos - pos0 mag"<< (pos-pos0).Mag() << "\n";
		//}

		GFDetPlane here,next;
		here.setO(pos);
		here.setNormal(mom);
		if(fVerbose) {std::cout<<"Original extrap direction:"<<"\n";mom.Print();std::cout<<"\n";}
		//if we successfully found a "first digi" use it as starting pos
		if (foundDigi >= 0 && fUseFirstDigiPos) {
			if(fVerbose) {
				std::cout<<"Moved plane from cluster position (r="<<here.getO().Perp()<<"): \n    " ;
				here.getO().Print();
				std::cout.flush();
				std::cout<<"\nto first Digi position (r="<<altFirstPos.Perp()<<"): \n    " ;
				altFirstPos.Print();
				std::cout.flush();
				std::cout<<"\n"<<"\n";
			}
			double pull = 0.99;
			if(altFirstPos.Perp() > 2*fRMin)  //first digi outside, track supposedly goes inwards
				pull = 1.01;
			here.setO(altFirstPos*pull);  //pull it inward a little
		}


		int counter(0);
		bool exc = false;

		std::set<unsigned int> usedDigis;  //bookkeeping for speedup

		std::cout<<"start stepping through track\n";
		while(true){ //stepalong until no digi is found between the planes
			bool abort = true;
			bool forceabort = false;
			double dir = 1.;
			if (fVerbose) std::cerr<<"at counter "<<counter<<"\n";
			if(counter++>0) here=next;//not the first time
			else if(here.getNormal() * startDir < 0) {
				//_DX*=-1;
				dir=-1.;
			}
			TVector3 destination = here.getO()+(here.getNormal()*_DX*dir);

			TVector3 poca,dirInPoca;
			double dist;
			try{
				theRep->extrapolateToPoint(destination,poca,dirInPoca);
				next.setO(poca);
				next.setNormal(dirInPoca);
			}
			catch(GFException& e){
				e.what();
				exc=true;
				break;
				//TODO: exception handling
			}

			//check if next lies out of the TPC volume
			//if (fVerbose) std::cerr<<"next ) perp: "<<next.getO().Perp()<<" rmax:"<<fRMax<<" rmin:"<<fRMin<<"\n";
			if(next.getO().Perp() > fRMax || next.getO().Perp() < fRMin) {
				if(fVerbose)
					std::cout<<"OUT OF CHAMBER  next.Perp(): <<"<<next.getO().Perp()<<"\n";
				TVector3 conn = next.getO() - here.getO();
				conn.SetMag(1.);
				double drn, drh, dr;
				dr = next.getO().Perp() - here.getO().Perp();
				double comp = fRMax;
				if(next.getO().Perp()<fRMin)
					comp = fRMin;
				drn = next.getO().Perp() - comp;
				double moveR = 1.1*fabs(drn)*_DX/fabs(dr);   //scaled to get safely inside the cage
				TVector3 newpos = next.getO() - (conn*moveR);
				next.setO(newpos);
				if(fVerbose)
					std::cout<<"AFTER SHIFT  next.Perp(): <<"<<next.getO().Perp()<<"\n";
				forceabort = true;
			}

			//now do the final dE/dx extrapolation
			try {
				dist = theRep->extrapolate(next,protState);
			}
			catch(GFException& e) {
				e.what();
				exc=true;
				break;
			}

			double thisDist = dist-totDist;
			totDist = dist;

			if (fVerbose) {
				std::cout<<"Extrapolation results :\n";poca.Print();
				std::cout<<"\n";dirInPoca.Print();
				std::cout<<"Extrap distance: "<<dist<<"\n";
			}

			TVector3 normHere = here.getNormal();
			TVector3 normNext = next.getNormal();

			double dE = 0.;

			double behindHere;
			double behindNext;


			if (fVerbose) {
				std::cerr<<"here:\n";
				here.Print();
				std::cerr<<"next:\n";
				next.Print();
			}

			std::map<unsigned int, TpcDigi*>::const_iterator cit;
			if(fVerbose) std::cout<<"start collecting digis in between ("<<digis.size()<<")\n";
			for(cit=digis.begin(); cit!=digis.end();cit++){
				if(usedDigis.count(cit->first))
					continue;
				TVector3 digiPos;
				try {
					TpcDigiMapper::getInstance()->map(cit->second,digiPos);
				}
				catch(const std::exception& ex) {
					std::cout<<ex.what()<<"\n";
					continue; //THIS WAS MISSING!
				}
				if(fIgnoreEdgeDigis)
					if(digiPos.Perp() < fRCutoff1 || digiPos.Perp() > fRCutoff2)
						continue;
				digiPos=TpcAlignmentManager::getInstance()->localToMaster("tpc",digiPos);

				behindHere = normHere * (here.dist(digiPos));
				behindNext = normNext * (next.dist(digiPos));

				behindHere *= 1./fabs(behindHere);
				behindNext *= 1./fabs(behindNext);


				double ok = behindHere + behindNext;    // 0 if digi between planes
				if (fVerbose) std::cout<<"JDIOSFSDF   OK:  "<<ok<<"\n";
				if (behindHere < 0.) {
					abort = false;  //there are still viable digis up fronts, keep running
					if(fabs(ok) < 0.5) {  //possibilities are actually -2, 0 ,2 within double precision...
						dE+=(cit->second)->amp();
						usedDigis.insert(cit->first);
					}
				}
			}

			if (fVerbose) std::cout<<"dE: "<<dE<<",  dist: "<<dist<<",  abort: "<<abort<<"\n";
			if(abort || forceabort)
				break;

			if(dE >= 0.)
				dedx.add(dE,thisDist);

			//MC PART ---------------------------------------------------

			dE = 0.;

			if(_idealdEdx) {
				for(unsigned int m=0;m<pointlist.size();++m){
					TVector3 pointPos;
					pointlist.at(m)->Position(pointPos);
					if((here.getO()-pointPos).Mag() > (_DX+2.)) continue;
					behindHere = normHere * (here.dist(pointPos));
					behindNext = normNext * (next.dist(pointPos));
					if(behindHere<0.){//in front of
						if(behindNext>=0.){
							dE+=pointlist.at(m)->GetEnergyLoss();
						}
					}
					else{
						if(behindNext<0.){//in front of
							dE+=pointlist.at(m)->GetEnergyLoss();
						}
					}
				}

				if(dE>0. && (!exc)){
					dedx_MC.add(dE*1E9,dist);
				}
			} // end if _idealdEdx
		}

		if(exc==false)
			new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);
		else
			new((*_dEdxOutArray)[itr]) TpcdEdx();

		if(_idealdEdx) {
			new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
		}
		//for(unsigned int i=0;i<digis.size();++i){
		//  delete digis[i];
		//}

	}
}

void 
TpcdEdxTask::WriteHistograms(const TString& filename){
	TFile* file = new TFile(filename,"RECREATE");
	file->mkdir("DEDX");
	file->cd("DEDX");

	_distHist->Write();
	delete _distHist;

	_dirHist->Write();
	delete _dirHist;

	file->Close();
	delete file;
}

void 
TpcdEdxTask::SetIgnoreEdgeDigis(double r1, double r2, Bool_t opt) {
	fRCutoff1 = r1;
	fRCutoff2 = r2;
	fIgnoreEdgeDigis = opt;
}



ClassImp(TpcdEdxTask)