#include "TpcClusterCorrector.h"
#include "TClonesArray.h"
#include "TpcPadPlane.h"
#include "TF1.h"

ClassImp(TpcClusterCorrector);

struct setXY
{
	setXY() : X(-1),Y(-1),Amp(0) {}
	double X;
	double Y;
	double Amp;
};
struct setT
{
	setT() : Amp(0),T(-1) {}
	double T;
	double Amp;
};

TpcClusterCorrector::TpcClusterCorrector(TpcPadPlane* plane, TClonesArray* digiArray)
:
		fparSetXY(false),
		fparSetT(false),
		fparSetXY2(false),
		fparSetT2(false),
		fpadplane(plane),
		fdigiArray(digiArray)

{
	for (int i=0;i<4;i++)
	{
		corrFuncXY.push_back(new TF1(Form("fitfuncXY_%d",i),
				"[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));

		corrFuncXY2.push_back(new TF1(Form("fitfuncXY2_%d",i),
				"[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));

	}

	corrFuncT  = new TF1("fitfuncT",
			"[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);
	corrFuncT2 = new TF1("fitfuncT2",
			"[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);

}
TpcClusterCorrector::TpcClusterCorrector()
:
		fparSetXY(false),
		fparSetT(false),
		fparSetXY2(false),
		fparSetT2(false),
		fpadplane(0),
		fdigiArray(0)

{
	for (int i=0;i<4;i++)
	{
		corrFuncXY.push_back(new TF1(Form("fitfuncXY_%d",i),
				"[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));

		corrFuncXY2.push_back(new TF1(Form("fitfuncXY2_%d",i),
				"[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));

	}

	corrFuncT  = new TF1("fitfuncT",
			"[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);
	corrFuncT2 = new TF1("fitfuncT",
			"[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);
}


TVector3 TpcClusterCorrector::ExecCorrection(TpcCluster* cl,bool doXY, bool doT)
{
	TVector3 newpos=cl->pos();

	if (doXY and not (fparSetXY or fparSetXY2))
	{
		std::cout<<"Cluster Correction in XY wanted but no parameters set!\n";
		return newpos;
	}
	if(doT and not (fparSetT or fparSetT2))
	{
		std::cout<<"Cluster Correction in T wanted but no parameters set!\n";
		return newpos;
	}

	const std::map<unsigned int, double> *digiMap=cl->getDigiMap();
	std::map<unsigned int, double>::const_iterator digiIt=digiMap->begin();
	std::map<unsigned int, double>::const_iterator digiEndIt=digiMap->end();
	std::map<unsigned int,setXY> collectionXY;
	std::map<double,setT> collectionT;
	//first of all collect all digis in sets. amplitudes of multiple used pads/zbins are added up.
	for(;digiIt!=digiEndIt;++digiIt)
	{
		TpcDigi* adigi=(TpcDigi*) fdigiArray->At((*digiIt).first);
		unsigned int padId=adigi->padId();
		collectionXY[padId].Amp+=adigi->amp();
		fpadplane->GetPadXY( padId,collectionXY[padId].X,collectionXY[padId].Y);
		collectionT[adigi->t()].Amp+=adigi->amp();
		collectionT[adigi->t()].T=adigi->t();
	}
	//check if correction is needed at all:
	if(collectionXY.size()>3 and collectionT.size()>3)
		return newpos;

	//look for the maximum amplitude in XY set
	std::map<unsigned int,setXY>::const_iterator collIt=collectionXY.begin();
	std::map<unsigned int,setXY>::const_iterator collEnd=collectionXY.end();
	setXY maxXY;
	setXY minXY;
	minXY.Amp=99999999;
	for (;collIt!=collEnd;++collIt)
	{
		double amp=(collIt->second).Amp;
		if(amp>maxXY.Amp)
			maxXY=collIt->second;
		if(amp<minXY.Amp)
			minXY=collIt->second;
	}

	//look for the maximum amplitude in T set
	std::map<double,setT>::const_iterator collItT=collectionT.begin();
	std::map<double,setT>::const_iterator collEndItT=collectionT.end();
	setT maxT;
	setT minT;
	minT.Amp=999999999;
	for(;collItT!=collEndItT;++collItT)
	{
		double amp=(collItT->second).Amp;
		if(amp>maxT.Amp)
			maxT=collItT->second;
		if(amp<minT.Amp)
			minT=collItT->second;
	}

	//now we calculate the direction and the amount of correction

	TVector3 axisXY(minXY.X-maxXY.X,minXY.Y-maxXY.Y,0);
	int rmag=floor( axisXY.Mag()*1000 + 0.5 );
	int ifunc=0;
	//std::cout<<"rmag="<<rmag<<"\n";
	switch (rmag)
	{
	case 260: ifunc=0; break;
	case 450: ifunc=1; break;
	case 520: ifunc=2; break;
	case 687: ifunc=3; break;
	default: ifunc=3; break;
	}
	//std::cout<<"ifunc="<<ifunc<<"\n";

	if(axisXY.Mag()!=0)
		axisXY.SetMag(1);
	double etaXY=minXY.Amp/(minXY.Amp+maxXY.Amp);

	if (collectionXY.size()==2 and doXY)
	{
		newpos-=axisXY*corrFuncXY[ifunc]->Eval(etaXY);
		//std::cout<<"shifted cluster position by "<<corrFuncXY[ifunc]->Eval(etaXY)<<"(ifunc:"<<ifunc<<", eta:"<<etaXY<<")"<<" along: \n";
		//corrFuncXY[ifunc]->Print();
		//for(int j=0;j<6;j++)
		//	std::cout<<"p"<<j<<" "<<corrFuncXY[ifunc]->GetParameter(j)<<"\n";
		//axisXY.Print();
	}
	else if(collectionXY.size()==3 and doXY)
	{
		newpos-=axisXY*corrFuncXY2[ifunc]->Eval(etaXY);
	}

	TVector3 axisT(0,0,minT.T-maxT.T);
	if (axisT.Mag()!=0)
		axisT.SetMag(1);
	double etaT=minT.Amp/(minT.Amp+maxT.Amp);
	if(collectionT.size()==2 and doT)
	{
		newpos-=axisT*corrFuncT->Eval(etaT);
		//std::cout<<"shifted cluster position by "<<etaT<<" along: \n";
		//axisT.Print();
	}
	else if (collectionT.size()==3 and doT)
	{
		newpos-=axisT*corrFuncT2->Eval(etaT);
	}

	return newpos;
}


void TpcClusterCorrector::SetParameterXY(double* parArr) {

	int ifunc=0;
	for(unsigned int ip=0; ip<4*6; ip++)
	{  //TODO: define this nicer
		ifunc=int(ip/6);
		std::cout<<"XY parameter: "<<ip%6<<"  "<<parArr[ip]<<" ifunc"<<ifunc<<"\n";
		corrFuncXY[ifunc]->SetParameter(ip%6, parArr[ip]);  //no range checks!

	}
	fparSetXY=true;
	}
void TpcClusterCorrector::SetParameterT(double* parArr) {
	for(unsigned int ip=0; ip<6; ip++)
	{  //TODO: define this nicer
		std::cout<<"T parameter: "<<ip<<"  "<<parArr[ip]<<"\n";
		corrFuncT->SetParameter(ip, parArr[ip]);  //no range checks!
	}
	fparSetT=true;
}
void TpcClusterCorrector::SetParameterXY2(double* parArr) {

	int ifunc=0;
	for(unsigned int ip=0; ip<4*6; ip++)
	{  //TODO: define this nicer
		std::cout<<"XY parameter: "<<ip<<"  "<<parArr[ip]<<" ifunc"<<ifunc<<"\n";
		ifunc=int(ip/6);
		corrFuncXY2[ifunc]->SetParameter(ip, parArr[ip]);  //no range checks!
	}
	fparSetXY2=true;
	}
void TpcClusterCorrector::SetParameterT2(double* parArr) {
	for(unsigned int ip=0; ip<6; ip++)
	{  //TODO: define this nicer
		std::cout<<"T parameter: "<<ip<<"  "<<parArr[ip]<<"\n";
		corrFuncT2->SetParameter(ip, parArr[ip]);  //no range checks!
	}
	fparSetT2=true;
}