//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Correction of cluster positions based on fit of 
//      aquired data
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer    TUM            (original author)
//
//
//-----------------------------------------------------------

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

#include "TpcClusterCorrectionTask.h"
#include "TF1.h"
#include "TpcPadPlane.h"
#include <TClonesArray.h>
ClassImp(TpcClusterCorrectionTask)


TpcClusterCorrectionTask::TpcClusterCorrectionTask()
: FairTask("Tpc Cluster Correction"),
  _persistence(false),
  _parSet(false),
  _useDevMap(false),
  _forceManDrift(false),
  _useDigis(false),
  _clusterBranchName("TpcCluster_raw"),
  _digiBranchName("TpcSample"),
  _clusterOutName("TpcCluster"),
  _devMapName(""),
  _driftVel(0),
  fgas(0),
  fpadplane(0),
  _clusterArray(0),
  _digiArray(0),
  _outArray(0),
  _devMap(0),
  fpar(0)
{
	corrFunc = new TF1("fitfunc",
			"[0]+[1]*x+[2]*TMath::Power(x,2)+[3]*TMath::Power(x,3)+[4]*TMath::Power(x,4)+[5]*TMath::Power(x,5)",
			0,1);
	fVerbose=false;
}

TpcClusterCorrectionTask::~TpcClusterCorrectionTask()
{
	delete corrFunc;
}



InitStatus
TpcClusterCorrectionTask::Init() {

	std::cout<<"TpcClusterCorrectionTask::Init() Initializing Cluster corrector\n";
	//Get ROOT Manager
	FairRootManager* ioman= FairRootManager::Instance();

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

	// Get input collection
	_clusterArray=(TClonesArray*) ioman->GetObject(_clusterBranchName);

	if(_clusterArray==0 and not _useDigis)
	{
		Error("TpcIdealTrackingTask::Init","Cluster-array not found!");
		return kERROR;
	}

	_digiArray=(TClonesArray*) ioman->GetObject(_digiBranchName);

	if(_digiArray==0)
	{
		Error("TpcIdealTrackingTask::Init","Digi-array not found!");
		return kERROR;
	}
	fpadplane = fpar->getPadPlane();
	fgas = fpar->getGas();
	//if(_forceManDrift)
//	_driftVel = fpar->getManDriftVel();
//
//	if(_driftVel>0){
//		_forceManDrift=true;
//	}else{
//		_driftVel = fgas->VDrift();
//		_forceManDrift=false;
//	}

	if (_useDevMap)
	{
		if (_devMapName=="")
		{
			std::cout<<"TpcClusterCorrectionTask::Init() Reading devmap from parbase\n";
			_devMapName= fpar->getClustCorrFile();
		}
		std::cout<<"TpcClusterCorrectionTask::Init() using devmap: "<<_devMapName.Data()<<"\n";
		_devMap =  new TpcDevmapCyl(_devMapName,kTRUE);
		//_devMap =  new TpcDevmapCyl(_devMapName,_driftVel);
		if (!_devMap->loaded())
		{
			Fatal("Init","Deviation Map not properly loaded");
			return kERROR;
		}
		_devMap->printc();
		if(_useDigis)
			std::cout<<"TpcClusterCorrectionTask::Init() Correcting Digis not Cluster\n";
	}

	// create and register output array
	_outArray = new TClonesArray("TpcCluster");
	ioman->Register(_clusterOutName,"Tpc",_outArray,_persistence);

	return kSUCCESS;
}

void
TpcClusterCorrectionTask::SetParContainers() {

	std::cout<<"TpcClusterFinderTask::SetParContainers"<<std::endl;
	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 TpcClusterCorrectionTask::Exec(Option_t* opt)
{
	if (_useDevMap)
	{
		if (fVerbose) std::cout<<"ClusterCorrectionTask::Exec using DevMap\n";
		if(_useDigis)
			ExecMapDigi(opt);
		else
			ExecMap(opt);
	}
	else
		ExecPar(opt);
}


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

	// Reset output Array
	if(_outArray==0)
		Fatal("TpcClusterCorrectionTask::Exec)","No OutputArray!");

	_outArray->Delete();

	if(!_parSet) {
		std::cerr<<"TpcClusterCorrectionTask::Exec() Correction parameters not properly set!"
				<<"... Aborting"<<std::endl;
		return;
	}

	//loop over clusters
	unsigned int nCl = _clusterArray->GetEntriesFast();
	for(unsigned int ic=0; ic<nCl; ic++) {
		TpcCluster* cl = (TpcCluster*) _clusterArray->At(ic);
		TpcCluster* cp = new ((*_outArray)[ic]) TpcCluster(*cl);
		TVector3 clpos = cp->pos();



		if(!(cp->size() == 2 && cl->nPad()==2))
			continue;
		//loop over digigs
		std::vector<TVector3*> positions;
		std::vector<double> amps;
		unsigned int nDigi = cp->nDigi();
		assert(nDigi==2);
		const std::map<unsigned int, double> *digiMap=cp->getDigiMap();
		std::map<unsigned int, double>::const_iterator digiIt=digiMap->begin();
		std::map<unsigned int, double>::const_iterator digiEndIt=digiMap->end();
		for(;digiIt!=digiEndIt;++digiIt){
			double x,y;
			TpcDigi* adigi=(TpcDigi*) _digiArray->At((*digiIt).first);
			fpadplane->GetPadXY( adigi->padId(),x,y);
			positions.push_back(new TVector3(x,y,0.)); //very inefficient, I know ...
			amps.push_back(adigi->amp());
		}
		//now construct cluster axis and correct
		TVector3 axis = *(positions[0]) - *(positions[1]);
		//define positive axis direction along positive x-axis
		if(axis.X()<0)
			axis*=(-1.);
		axis.SetMag(1.); //normalize

		//calculate eta (fraction of left amplitude and total amplitude)
		double smallAmp = amps[0];
		if(positions[1]->X() < positions[0]->X())
			smallAmp = amps[1];
		double eta = smallAmp/(amps[0]+amps[1]);

		double shift = corrFunc->Eval(eta);
		TVector3 newPos = clpos+(shift*axis);
		cp->SetPos(newPos);

		std::cout<<"    shifted cluster position by "<<shift<<" along cluster axis"
				<<std::endl;
		for(unsigned int i=0; i<positions.size(); i++)
			delete positions[i];

	} //end loop over clusters
	return;

}

void
TpcClusterCorrectionTask::SetParameters(double* parArr) {
	for(unsigned int ip=0; ip<6; ip++) {  //TODO: define this nicer
		corrFunc->SetParameter(ip, parArr[ip]);  //no range checks!
	}
	_parSet=true;
}

void TpcClusterCorrectionTask::ExecMap(Option_t * opt)
{
	// Reset output Array
	if(_outArray==0)
		Fatal("TpcClusterCorrectionTask::Exec)","No OutputArray!");

	_outArray->Delete();

	//loop over clusters
	unsigned int nCl = _clusterArray->GetEntriesFast();
	for(unsigned int ic=0; ic<nCl; ic++)
	{
		TpcCluster* cl = (TpcCluster*) _clusterArray->At(ic);
		TpcCluster* cp = new ((*_outArray)[ic]) TpcCluster(*cl);
		TVector3 clpos = cp->pos();
		TVector3 corrval=_devMap->value(clpos);
		//corrval.SetZ(0);
		TVector3 newPos=clpos-corrval;
		if(fVerbose)
		{
			std::cout<<"ClusterCorrection Task::ExecMap: Old ("<<clpos.X()<<","<<clpos.Y()<<","<<clpos.Z()<<") New: ("<<newPos.X()<<","<<newPos.Y()<<","<<newPos.Z()<<")\n";
			std::cout<<"ClusterCorrection Task::ExecMap: correction: ("<<corrval.X()<<","<<corrval.Y()<<","<<corrval.Z()<<")\n";
		}


		cp->SetPos(newPos);
		cp->SetFieldCorr(corrval);
	}
}

void TpcClusterCorrectionTask::ExecMapDigi(Option_t * opt)
{
	// Reset output Array
	if(_outArray==0)
		Fatal("TpcClusterCorrectionTask::Exec)","No OutputArray!");

	_outArray->Delete();
	//loop over digis
	unsigned int nCl = _digiArray->GetEntriesFast();
	for(unsigned int ic=0; ic<nCl; ic++)
	{
		TpcDigi* digi = (TpcDigi*) _digiArray->At(ic);
		TVector3 digiPos;
		TpcDigiMapper::getInstance()->map(digi,digiPos);
		TVector3 corrval=_devMap->value(digiPos);
		TVector3 newPos=digiPos-corrval;
		//std::cout<<"old Pos: ("<<digiPos.X()<<","<<digiPos.Y()<<","<<digiPos.Z()<<") old Tick:"<<digi->t()<<"\n";
		int newID=fpadplane->GetPadList(newPos.X(),newPos.Y(),0)[0];
		digi->padId(newID);
		double newtick=TpcDigiMapper::getInstance()->tick_from_z(newPos.Z());
		digi->t(newtick);
		//std::cout<<"new Pos: ("<<newPos.X()<<","<<newPos.Y()<<","<<newPos.Z()<<") new Tick:"<<digi->t()<<"\n";
	}
}