#include "TpcClusterErrorScaler.h"
#include "TpcDigiPar.h"
#include "TpcDigiMapper.h"
#include "FairRootManager.h"
#include "TpcCluster.h"

#include "FairRun.h"
#include "FairRuntimeDb.h"

#include "TString.h"
#include "TMath.h"

#include <iostream>
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <string>
#include "TMatrixDSymEigen.h"


TpcClusterErrorScaler* TpcClusterErrorScaler::finstance=NULL;

TpcClusterErrorScaler::~TpcClusterErrorScaler()
{
	// do not delete the members! they might not belong to you!!!
}

TpcClusterErrorScaler::TpcClusterErrorScaler()
:fparFile(""),
 fInit(false),
 fNoFairRun(false),
 ffunc(56),
 dTheta(0),
 nTheta(0),
 dZ(0),
 nZ(0),
 fdx(0),
 fdy(0),
 fdz(0),
 fwx(0),
 fwy(0),
 fdecayX(0),
 fdecayY(0),
 fdecayZ(0),
 fDt(1),
 fDl(1),
 fManualDiffusion(kFALSE),
 fClusterArray(0),
 fClusterBranchName("TpcCluster")
{
	fdigiPar=NULL;
	//  fpar={1,0,0,0,1,0,0,0,1,0,0,0};
}

void TpcClusterErrorScaler::init(TString filename="",Int_t func=56)
{

	if(fInit) { //Protection against re-initialization
		std::string msg("TpcClusterErrorScaler has already been initialized!");
		std::runtime_error err(msg);
		throw err;
	}

	ffunc=func;

	if (filename!="")
		fparFile=filename;

	//std::cout<<"TpcClusterErrorScaler: fparfile="<<fparFile<<"   ffunc="<<ffunc<<"\n";
	if (fparFile!="")
	{
		if(ffunc==56)
		{
			readLookupFile();
		}
		else
		{
			readParFile();
		}
	}
	else
	{
		//std::string msg("Trying to use TpcClusterErrorScaler without parameter file! -> no scaling\n");
		//std::runtime_error err(msg);
		//throw err;
		std::cout<<"Trying to use TpcClusterErrorScaler without parameter file! -> no scaling\n";
		ffunc=-1;
	}

	std::cout<<"TpcClusterErrorScaler::SetParContainers"<<std::endl;
	std::cout.flush();

	if(not fNoFairRun)
	{
		// 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
		fdigiPar= (TpcDigiPar*) db->getContainer("TpcDigiPar");
		if (! fdigiPar ) Fatal("SetParContainers", "TpcDigiPar not found");

		//Get ROOT Manager ------------------
		FairRootManager* ioman= FairRootManager::Instance();
		if(ioman==0)
		{
			Error("TpcClusterErrorScaler::Init","RootManager not instantiated!");
		}

		fClusterArray=(TClonesArray*) ioman->GetObject(fClusterBranchName);
		if(fClusterArray==0)
		{
			Error("TpcClusterErrorScaler::Init","TpcCluster-Array not found!");
		}
		if (fparFile=="")
		{
			fparFile=fdigiPar->getClusterErrorFile();
			if (fparFile!="")
			{
				std::cout<<"TpcClusterErrorScaler::Init: never mind the error about missing parameter file, it was given in the run-parfile. switching back to scaling "<<func<<"\n";
				std::cout<<"TpcClusterErrorScaler::Init: ClusterErrorParfile="<<fparFile<<"\n";
				ffunc=func;
				if(ffunc==56)
				{
					readLookupFile();
				}
				else
				{
					readParFile();
				}
			}
		}
	}

	fInit=true;
	std::cout<<"initialized error scaler with: ffunc="<<ffunc<<"\n";
}

void TpcClusterErrorScaler::postinit()
{
	if(fNoFairRun)
		return;
	if(fdx==0 or fdy==0 or fdz==0 or fdecayX==0 or fwx==0 or fwy==0)
	{
		double dx, dy;
		TpcDigiMapper::getInstance()->padsize(0,dx,dy);
		fwx=(TpcDigiMapper::getInstance()->getPadPlane()->GetPad(0)->GetWeight(0.99,1.,0));
		fwy=(TpcDigiMapper::getInstance()->getPadPlane()->GetPad(0)->GetWeight(0.99,0,1.));

		fdx=dx/2;
		fdy=dy/2*(1+sin(30*TMath::DegToRad()));

		//this block is to define the z jitter
		McId dummyID(1,1,0);
		McIdCollection dummyColl;
		dummyColl.AddID(dummyID);

		TVector3 zDiff1,zDiff2;
		TpcDigi zDiffDigi1(1,1,1,dummyColl),zDiffDigi2(1,2,1,dummyColl);
		try
		{
			TpcDigiMapper::getInstance()->map(&zDiffDigi1,zDiff1);
			TpcDigiMapper::getInstance()->map(&zDiffDigi2,zDiff2);
		}
		catch(const std::exception& ex) {
			std::cout<<ex.what()<<std::endl;
		}
		fdz = zDiff2.z() - zDiff1.z();

		double Dx=fwx;
		double Dy=fwy;
		double Dz=fdz;
	}
}

Double_t  TpcClusterErrorScaler::scaleFunc(Double_t driftl, Int_t i,TVector3 ax,TVector3 diff)
{
	Double_t scale=1.;
	switch (ffunc)
	{
	case 0:
		scale=scaleFunc_phys(driftl,i);
		break;
	case 1:
		scale=scaleFunc_pol(driftl,i);
		break;
	case 2:
		scale=scaleFunc_diff(driftl,i,ax);
		break;
	case 20:
		scale=scaleFunc_diff(driftl,i,ax)*scaleFunc_phys(driftl,i);
		break;
	case 3:
		scale=scaleFunc_mean(driftl,0);
		break;
	case 4:
		scale=scaleFunc_mean(driftl,1);
		break;
	default:
		scale=1;
		break;
	}
	return scale;
}

Double_t TpcClusterErrorScaler::scaleFunc(Double_t driftl, Double_t theta, Int_t i, bool mirr)
{
	Double_t scale=1.;

	if (theta!=theta)
		return scale;

	if (ffunc==56)
	{
		if(mirr)
			scale=scaleFunc_lookup2(driftl,theta,i);
		else
			scale=scaleFunc_lookup(driftl,theta,i);
	}
	else if (ffunc==57)
		//scale=5;
		scale=scaleFunc_pol(driftl,i);
	//std::cout<<"last time in scaler:"<<scale<<"\n";

	return scale;
}

Double_t TpcClusterErrorScaler::scaleFunc_mean(Double_t driftl,Int_t opt)
{
	Double_t val=0;

	for(int i=0;i<2;i++)
	{
		if (opt==0)
			val+=scaleFunc_pol(driftl,i);
		else if (opt==1)
			val+=scaleFunc_phys(driftl,i);
	}
	val/=2;

	return val;
}

Double_t TpcClusterErrorScaler::scaleFunc_phys(Double_t driftl, Int_t i)
{
	Double_t val=1.;

	if (driftl>=0)
		val=fpar[i][0]+fpar[i][1]*sqrt(driftl)+fpar[i][2]*exp(-pow(driftl,fpar[i][3]));
	else
		val=1.;

	return fabs(val);
}

Double_t TpcClusterErrorScaler::scaleFunc_pol(Double_t driftl, Int_t i)
{
	Double_t val=1.;
	if (driftl>=0)
	{
		val=fpar[i][0]
		            +fpar[i][1]*driftl
		            +fpar[i][2]*pow(driftl,2)
		            +fpar[i][3]*pow(driftl,3)
		            +fpar[i][4]*pow(driftl,4)
		            +fpar[i][5]*pow(driftl,5)
		            +fpar[i][6]*pow(driftl,6);
	}
	else
		val=1.;

	return val;
}

Double_t TpcClusterErrorScaler::scaleFunc_diff(Double_t driftl, Int_t i,TVector3 ax)
{
	double Dl = TpcDigiMapper::getInstance()->getGas()->Dl();
	double Dt = TpcDigiMapper::getInstance()->getGas()->Dt();

	TVector3 VDiff;
	VDiff.SetXYZ(Dl,Dl,Dt);
	VDiff*=TMath::Sqrt(driftl);

	Double_t val=1.;

	val=VDiff*ax/ax.Mag();

	return val;
}

Double_t TpcClusterErrorScaler::scaleFunc_lookup2(Double_t driftl, Double_t theta, Int_t i)
{
	Double_t val=1;
	val =scaleFunc_lookup(driftl,theta,i);
	val+=scaleFunc_lookup(driftl,180-theta,i);

	val/=2.;

	return val;
}

Double_t TpcClusterErrorScaler::scaleFunc_lookup(Double_t driftl, Double_t theta, Int_t i)
{
	Double_t val=1;

	double zbin=driftl/dZ;
	double thetabin=theta/dTheta;

	int lz=int(zbin);
	int ltheta=int(thetabin);
	int hz=lz+1;
	int htheta=ltheta+1;

	if (hz>=nZ)
		if(htheta>=nTheta)
			val=flookupTable[i].back().back();
		else
			val=linInterpol(flookupTable[i][ltheta].back(),flookupTable[i][htheta].back(),thetabin-ltheta);
	else
		if(htheta>=nTheta)
			val=linInterpol(flookupTable[i].back()[lz],flookupTable[i].back()[hz],zbin-lz);
		else
		{
			val=biLinInterpol(flookupTable[i][ltheta][lz],flookupTable[i][htheta][lz],flookupTable[i][ltheta][hz],flookupTable[i][htheta][hz],thetabin-ltheta,zbin-lz);
		}
	return val;
}

Double_t TpcClusterErrorScaler::linInterpol(double v1,double v2, double x)
{
	double retval=0;
	retval=v1*(1-x)+(x)*v2;
	//std::cout<<"lin interpolation with:"<<v1<<" "<<v2<<" "<<x<<" "<<retval<<"\n";
	return retval;
}

Double_t TpcClusterErrorScaler::biLinInterpol(double v1, double v2, double v3, double v4, double x, double y)
{
	double retval=0;
	double inter1=linInterpol(v1,v2,x);
	double inter2=linInterpol(v3,v4,x);

	retval=linInterpol(inter1,inter2,y);

	return retval;
}

void TpcClusterErrorScaler::readParFile()
{
	ifstream input(fparFile.Data());
	if(!input)
	{
		//std::string msg("TpcClusterErrorScaler: could not open par file!");
		//std::runtime_error err(msg);
		//throw err;
		std::cout<<"TpcClusterErrorScaler: could not open par file!\n";
		ffunc=-1;
		return;
	}

	std::string line;
	int counter=-1;
	std::cout<<"TpcClusterErrorScaler: reading file "<<fparFile.Data()<<"\n";
	while(getline(input,line))
	{
		counter++;
		if (counter>2)
		{
			counter--;
			break;
		}
		std::istringstream lineStream(line);
		double a,b,c,d,e,f,g,h;
		a=1;
		b=0;
		c=0;
		d=0;
		e=0;
		f=0;
		g=0;
		h=0;
		if (ffunc==0 || ffunc==20 || ffunc==4)
		{
			if (!(lineStream >> a >> b >> c >> d ))
			{
				std::string msg("TpcClusterErrorScaler: error while reading parfile!");
				std::runtime_error err(msg);
				throw err;
			}
		}
		else if (ffunc==1 || ffunc==3 || ffunc==57)
		{
			if (!(lineStream >>a>>b>>c>>d>>e>>f>>g))
			{
				std::string msg("TpcClusterErrorScaler: error while reading parfile!");
				std::runtime_error err(msg);
				throw err;
			}
		}
		else if (ffunc==61)
		{
			if (!(lineStream >> a >> b >> c >> d >> e >> f >> g >> h))
			{
				std::string msg("TpcClusterErrorScaler: error while reading parameter file!");
				std::runtime_error err(msg);
				throw err;
			}
		}
		fpar[counter][0]=a;
		fpar[counter][1]=b;
		fpar[counter][2]=c;
		fpar[counter][3]=d;
		fpar[counter][4]=e;
		fpar[counter][5]=f;
		fpar[counter][6]=g;
		fpar[counter][7]=h;

		std::cout<<"found parameter: "<<a<<" "<<b<<" "<<c<<" "<<d<<" "<<e<<" "<<f<<" "<<g<<" "<<h<<"\n";

	}
	std::cout<<"TpcClsuterErrorScaler: read file with "<<counter+1<<" entries\n";
	if (counter!=2 and ffunc!=61)
	{
		std::string msg("TpcClusterErrorScaler: not enough entries in parfile!\n");
		std::runtime_error err(msg);
		throw err;
	}
}

void TpcClusterErrorScaler::readLookupFile()
{
	ifstream input(fparFile.Data());
	if(!input)
	{
		std::cout<<"TpcClusterErrorScaler: could not open par file!\n";
		ffunc=-1;
		return;
	}

	std::string line;
	int counter=-1;

	getline(input,line);
	std::istringstream lineStream_first(line);
	//double nTheta, nZ, dTheta, dZ;
	if(!(lineStream_first>>nTheta>>dTheta>>nZ>>dZ))
	{
		std::string msg("TpcClusterErrorScaler: error while reading lookup file header line!");
		std::runtime_error err(msg);
		throw err;
	}


	std::vector< std::vector< std::vector<double> > > A(2,std::vector<std::vector<double> >(nTheta,std::vector<double>(nZ)));

	flookupTable=A;

	std::cout<<"TpcClusterErrorScaler: reading file "<<fparFile.Data()<<"\n";
	while(getline(input,line))
	{
		counter++;

		std::istringstream lineStream(line);
		double axis;
		double theta;
		double zpos;
		double scale;

		if(!(lineStream>>axis>>theta>>zpos>>scale))
		{
			std::string msg("TpcClusterErrorScaler: error while reading lookup file!");
			std::runtime_error err(msg);
			throw err;
		}
		int theta_bin=int(theta/dTheta);
		int z_bin=int(zpos/dZ);

		flookupTable[int(axis)][theta_bin][z_bin]=scale;

	}
	std::cout<<"TpcClusterErrorScaler: successfully read lookuptable\n";
}

void TpcClusterErrorScaler::paramVariance(double& varX, double& varY, double& varZ, Double_t driftl, TVector3 dir)
{
	postinit();
	if (!ffunc>=60)
	{
		std::cout<<"parametrized Variance requested in wrong scaling mode!\n";
		varX=varY=varZ=0;
		return;
	}

	double Dl = TpcDigiMapper::getInstance()->getGas()->Dl();
	double Dt = TpcDigiMapper::getInstance()->getGas()->Dt();
	double theta=dir.Theta();
	double phi=dir.Phi();

	/*
	varX=fpar[0][0] + fpar[0][1]*Dt*Dt*driftl + (4*fwx*fwx)/( 12 )*TMath::Exp(-fpar[0][2]*driftl);
	varY=fpar[1][0] + fpar[1][1]*Dt*Dt*driftl + (4*fwy*fwy)/( 12 )*TMath::Exp(-fpar[1][2]*driftl);
	varZ=fpar[2][0] + fpar[2][1]*Dl*Dl*driftl + ( fdz*fdz )/( 12 )*TMath::Exp(-fpar[2][2]*driftl);
	 */

	varX=fwx/3.+Dt*Dt*driftl;
	varY=fwy/3.+Dt*Dt*driftl;
	varZ=fdz/6.+Dl*Dl*driftl;
}

TMatrixDSym TpcClusterErrorScaler::getErrorCov(double driftl, TVector3 dir)
{
	postinit();
	TMatrixDSym errorCov(3);
	double Dl;
	double Dt;


	if(fManualDiffusion)
	{
		Dl=fDl;
		Dt=fDt;
	}
	else
	{
		Dl = TpcDigiMapper::getInstance()->getGas()->Dl();
		Dt = TpcDigiMapper::getInstance()->getGas()->Dt();
	}

	double theta = dir.Theta();
	double ct=TMath::Cos(theta);
	double st=TMath::Sin(theta);

	//simple linear offset
	TMatrixDSym offsetCov(3);
	offsetCov[0][0]=fpar[0][0];
	offsetCov[1][1]=fpar[1][0];
	offsetCov[2][2]=fpar[2][0];

	TMatrixDSym padCov(3);
	padCov[0][0]=4*(fwx*fwx); //*4 since fwx is half pad size
	padCov[1][1]=4*(fwy*fwy); //*4 since fwy is half pad size
	padCov[2][2]=1*(fdz*fdz); //*1 since fdz is full z bin
	padCov*=fpar[0][0];

	//diffusion driven spread of the electrons
	TMatrixDSym diffCov(3);
	diffCov[0][0]+= fpar[0][1]*Dt*Dt*driftl;
	diffCov[1][1]+= fpar[1][1]*Dt*Dt*driftl;
	diffCov[2][2]+= fpar[2][1]*Dl*Dl*driftl;


	//term linear in Z:
	TMatrixDSym linCov(3);
	linCov[0][0]+=fpar[0][2]*pow(driftl,2);
	linCov[1][1]+=fpar[1][2]*pow(driftl,2);
	linCov[2][2]+=fpar[2][2]*pow(driftl,2);

	//exponential with driftpos decaying part due to finite pad size
	double dx=2*fwx;
	double dy=2*fwy;
	double dz=fdz;
	double padErrX=dx*dx/( 12 );
	double padErrY=dy*dy/( 12 );
	double padErrZ=dz*dz/( 12 );
	TMatrixDSym padDecayCov(3);
	padDecayCov[0][0]+=fpar[0][3]*TMath::Exp(-fpar[0][2]*driftl)*pow(st,2);
	padDecayCov[1][1]+=fpar[1][3]*TMath::Exp(-fpar[1][2]*driftl)*pow(st,2);
	padDecayCov[2][2]+=fpar[2][3]*TMath::Exp(-fpar[2][2]*driftl)*pow(st,2);


	//various angle models
	TMatrixDSym angleCov(3);

	angleCov[0][0]=fpar[0][4]*pow(ct,2);
	angleCov[1][1]=fpar[1][4]*pow(ct,2);
	angleCov[2][2]=fpar[2][4]*pow(st,2);

	TMatrixDSym angleCov2(3);
	angleCov2[0][0]=fpar[0][5]*pow(ct,4);
	angleCov2[1][1]=fpar[1][5]*pow(ct,4);
	angleCov2[2][2]=fpar[2][5]*pow(st,4);

	//mixed term diffusion and angle
	TMatrixDSym mixedCov(3);
	mixedCov[0][0]+=fpar[0][6]*driftl*pow(ct,2);
	mixedCov[1][1]+=fpar[1][6]*driftl*pow(ct,2);
	mixedCov[2][2]+=fpar[2][6]*driftl*pow(st,2);


	errorCov+=offsetCov;
	errorCov+=padDecayCov;
	errorCov+=diffCov;
	//errorCov+=linCov;
	errorCov+=angleCov;
	errorCov+=angleCov2;
	errorCov+=mixedCov;
	//errorCov+=padCov;

	//errorCov.Invert();
	//errorCov=errorCov.SimilarityT(fPlaneM);
	//errorCov.Invert();
	//std::cout<<" zpos:"<<driftl<<" theta:"<<dir.Theta()<<"\n";
	//errorCov.Print();

	return errorCov;
}

TMatrixDSym TpcClusterErrorScaler::getErrorCovOnPlane(int clusterID, TMatrixD plane, TMatrixDSym shapeCovPlane, double _zpos)
{
	TMatrixDSym planeCov(2);
	TMatrixDSym shapeOnPlane(2);

	double S1=-1;
	double S2=-1;

	double zpos;
	if (clusterID==-1)
	{
		shapeOnPlane=TMatrixDSym(2,shapeCovPlane.GetMatrixArray());
		zpos=_zpos;
	}
	else
	{
		TpcCluster* clus=(TpcCluster*)fClusterArray->At(clusterID);
		shapeOnPlane=TMatrixDSym(2,clus->shapeCovPlane().GetMatrixArray());
		zpos=clus->pos().Z();
		//if (clus->isGeom().Y())
		//	S1=clus->minSigma().Y()/12.;
		//if (clus->isGeom().Z())
		//	S2=clus->minSigma().Z()/12.;
	}
	TVector3 U; U.SetXYZ(plane[0][0],plane[1][0],plane[2][0]);
	TVector3 V; V.SetXYZ(plane[0][1],plane[1][1],plane[2][1]);

	double theta = U.Cross(V).Theta();
	double ct=TMath::Cos(theta);
	double st=TMath::Sin(theta);

	TMatrixDSymEigen eigen(shapeOnPlane);
	TMatrixD ev=eigen.GetEigenVectors();
	TMatrixD ev_i(2,2,ev.GetMatrixArray());
	ev_i.Invert();
	TMatrixD S(2,2);

	if (S1!=-1)
		S[0][0]=S1;
	else
		S[0][0]=fpar[0][0]+fpar[0][1]*zpos+fpar[0][3]*TMath::Exp(-fpar[0][2]*zpos) + fpar[0][4]*pow(ct,2);

	if(S2!=-1)
		S[1][1]=S2;
	else
		S[1][1]=fpar[2][0]+fpar[2][1]*zpos+fpar[2][3]*TMath::Exp(-fpar[2][2]*zpos) + fpar[2][4]*pow(ct,2);

	planeCov=TMatrixDSym(2, (ev*S*ev_i).GetMatrixArray());

	return planeCov;
}

double TpcClusterErrorScaler::paramAmp()
{
	return 1;//fpar[0][4]*400/fdigiPar->getGain();
}

void TpcClusterErrorScaler::setParameter(int parset, int parnum, double par)
{
	if(parset>2 or parnum>6)
	{
		std::cout<<"more paramter than possible\n";
		return;
	}
	fpar[parset][parnum]=par;
}