//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Correction of CdcSpacepoints
//
//
// Environment:
//      Software NOT developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sverre Doerheim    TUM            (original author)
//
//
//-----------------------------------------------------------

#include "CdcCorrectionTask.h"

#include "TClonesArray.h"
#include "TVector3.h"
#include "TVectorD.h"
#include "TH2D.h"
#include "TH1D.h"
#include "TFile.h"

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

#include "TpcDigiPar.h"
#include "TpcSPHit.h"
#include "TpcAlignmentManager.h"

#include "GFTrack.h"
#include "RKTrackRep.h"
#include "GFException.h"

#include "CdcSpacepoint.h"
#include "CdcHit.h"

#include <string>
#include <sstream>

CdcCorrectionTask::CdcCorrectionTask()
: FairTask("Cdc Correction task", 0),
  fcopy(false),
  fCdcHitBranchName("CdcHit"),
  fCdcHitOutBranchName(""),
  fpersistence(false),
  fCorrectionFactor(0.2),
  fCorrFileName(""),
  fCdcVdr(0),
  fCdcLorentz(0),
  fCdcT0(0),
  fCdcT0Corr(0),
  fFlightTimeCorr(false)
{
}

CdcCorrectionTask::~CdcCorrectionTask() {
	;
}



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

	// get input arrays
	fCdcHitArray = (TClonesArray*)ioman->GetObject(fCdcHitBranchName);
	if (fCdcHitArray==0)
		Fatal("CdcCorrectionTask::Init", "Input array not found! Aborting");
	if(fcopy) {
		fCdcHitOutArray = new TClonesArray("CdcHit");
		ioman->Register(fCdcHitOutBranchName,"tpc",fCdcHitOutArray,fpersistence);
	}

	for(unsigned int sectorId=1; sectorId<=16;++sectorId){
		TVector3 testVec=TVector3(50,0,0);
		TVector3 testVec2=TVector3(30,0,0);
		testVec2.SetPhi(-8*TMath::DegToRad());
		TVector3 innerPos=testVec-testVec2;
		TVector3 outerPos=testVec+testVec2;

		innerPos.SetPhi(innerPos.Phi()+(11.25+22.5*(sectorId-1))*TMath::DegToRad()); //fcgthit.f
		outerPos.SetPhi(outerPos.Phi()+(11.25+22.5*(sectorId-1))*TMath::DegToRad()); //fcgthit.f

		TVector3 corrVec=outerPos-innerPos;
		if(fVerbose)std::cout<<"cdcCorrectionTask: scaling corrVec by "<<fCorrectionFactor<<"\n";
		corrVec.SetMag(fCorrectionFactor);
		corrVec.RotateZ(-TMath::PiOver2());
		correctionVectors[sectorId]=corrVec;
		negCorrectionVectors[sectorId]=-corrVec;
		TVector3 driftVec=corrVec;
		if(fVerbose)std::cout<<"cdcCorrectionTask: scaling driftVec by "<<-fCdcVdr<<"\n";
		driftVec.SetMag(-fCdcVdr);
		driftVec.RotateZ(fCdcLorentz*TMath::DegToRad());
		driftVectors[sectorId]=driftVec;
		negDriftVectors[sectorId]=-driftVec;
	}

	if(readCorrFile(fCorrFileName)) {
		Warning("CdcCorrectionTask::Init",
				"Problems encountered during reading of correction file");
	}
	else {
		std::cout<<"CdcCorrectionTask::Init() finished without errors"<<std::endl;
	}

	return kSUCCESS;
}

void CdcCorrectionTask::Exec(Option_t* opt) {
	if(fVerbose) std::cout << "CdcCorrectionTask::Exec()\n";
	// Reset output Arrays

	int nCdcHits=fCdcHitArray->GetEntriesFast();
	for(int iHit=0;iHit<nCdcHits;++iHit){

		CdcHit* cdcHit=(CdcHit*) fCdcHitArray->At(iHit);
		int wireId=cdcHit->GetProperWireID();
		int sectorId=wireId%100;
		CdcHit* hitOut=cdcHit;
		if(fcopy){
			hitOut=new ((*fCdcHitOutArray)[fCdcHitOutArray->GetEntriesFast()])CdcHit(*cdcHit);
		}
		double driftLength=hitOut->GetHitDrift();
		TVector3 wirePlaneCor;
		if(correctionVectors.count(sectorId)){
			wirePlaneCor=correctionVectors[sectorId];
		}else{
			std::cerr<<"CdcCorrectionTask:Exec () sector id not found: "<<sectorId<<std::endl;
			throw;
		}
		TVector3 drift;
		if(fFlightTimeCorr){

			if(driftVectors.count(sectorId)){

				if(driftLength>0){
					drift=driftVectors[sectorId];
				}else{
					drift=negDriftVectors[sectorId];
				}

			}else{
				std::cerr<<"CdcCorrectionTask:Exec () sector id not found: "<<sectorId<<std::endl;
				throw;
			}


		}

		TVector3 t0Corr;
		if(fabs(fCdcT0Corr)!=0){
			if(driftVectors.count(sectorId)){
				t0Corr=driftVectors[sectorId];
			}else{
				std::cerr<<"CdcCorrectionTask:Exec () sector id not found: "<<sectorId<<std::endl;
				throw;
			}
			if(fVerbose)std::cout<<"cdcCorrectiontask: scaling t0Corr by "<<fCdcVdr*fCdcT0Corr<<"\n";
			t0Corr.SetMag(fCdcVdr*fCdcT0Corr);
		}

		TVector3 pos=hitOut->GetHitPos();
		//if(fVerbose)std::cout<<"before ("<<pos.X()<<","<<pos.Y()<<","<<pos.Z()<<")\n";
		pos=pos+wirePlaneCor;

		//z-correction:
		double zCorr = 0.;
		std::map<int,double>::const_iterator it = fCorrectionsZ.find(wireId);
		if(it!=fCorrectionsZ.end()) {
			zCorr = it->second;
			//std::cout<<"*($*#&*$&#$&    correcting wireHit by dz="<<zCorr<<std::endl;
		}

		if(fabs(fCdcT0Corr)!=0){
			if(driftLength>0){
				pos=pos+t0Corr;
			}else{
				pos=pos-t0Corr;
			}
		}
		//if(fVerbose)std::cout<<"after ("<<pos.X()<<","<<pos.Y()<<","<<pos.Z()<<")\n";
		hitOut->SetHitPos(pos.x(),pos.y(),pos.z()-zCorr);

		hitOut->SetT0(fCdcT0);
		hitOut->SetVdrift(fCdcVdr);
		hitOut->SetLorentzAngle(fCdcLorentz);
		if(fFlightTimeCorr)
		{
			if(fVerbose)std::cout<<"CdcCorrectionTask: flight time corr: ("<<drift.X()<<","<<drift.Y()<<","<<drift.Z()<<")\n";
			hitOut->SetDriftVector(drift);
		}
	}
}

TVector3 CdcCorrectionTask::GetCorrectionVector(int sector)
{
	if(correctionVectors.count(sector)){
		return correctionVectors[sector];
	}
	return TVector3(0,0,0);
}


void 
CdcCorrectionTask::SetParContainers() {
	std::cout << "CdcCorrectionTask::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");
}


int
CdcCorrectionTask::readCorrFile(const TString& cf) {
	std::cout<<"CdcCorrectionTask::readCorrFile() reading file "<<cf.Data()<<std::endl;
	TFile infile(fCorrFileName, "read");
	if(infile.IsZombie()) {
		Warning("CdcCorrectionTask::readCorrFile",
				"Correction file not found, no corrections will be performed");
		return 1;
	}


	unsigned int nkeys = infile.GetListOfKeys()->GetSize();
	//pattern is "WireOffsets_wireID", where whireID is secID + 100*wireID
	TString comp("wireOffset");
	for(unsigned int k=0; k<nkeys; k++) {
		TString name = infile.GetListOfKeys()->At(k)->GetName();
		if(name.Contains(comp)) {
			TH1D* offs = (TH1D*) infile.Get(name);
			std::string str_name(name.Data());
			int lo = str_name.find_last_of("_");
			std::string str_wid = str_name.substr(lo+1);
			int wid;
			std::stringstream convert(str_wid);
			if(!(convert >> wid)) {
				Fatal("CdcCorrectionTask::readCorrFile",
						"Could not convert identifier to integer! Aborting");
			}

			if(fCorrectionsZ.count(wid)) {
				Fatal("CdcCorrectionTask::readCorrFile",
						"WireIDs seem not to be unique in input! Aborting");
			}
			//fCorrectionsZ[wid] = offs->GetMean();
			TH1D* rebin = (TH1D*) offs->Rebin(5);
			int maxBin = rebin->GetMaximumBin();
			double maxPos = rebin->GetBinCenter(maxBin);
			fCorrectionsZ[wid] = maxPos;
			delete rebin;
			//std::cout<<"   setting mean to "<<offs->GetMean()<<" for wire "<<wid<<std::endl;

		}
	}

	infile.Close();
	return 0;

}

double CdcCorrectionTask::GetCdcT0Corr() const
{
	return fCdcT0Corr;
}

void CdcCorrectionTask::SetCdcT0Corr(double value)
{
	fCdcT0Corr = value;
}

double CdcCorrectionTask::GetCdcT0() const
{
	return fCdcT0;
}

void CdcCorrectionTask::SetCdcT0(double value)
{
	fCdcT0 = value;
}

double CdcCorrectionTask::GetCdcLorentz() const
{
	return fCdcLorentz;
}

void CdcCorrectionTask::SetCdcLorentz(double value)
{
	fCdcLorentz = value;
}

double CdcCorrectionTask::GetCdcVdr() const
{
	return fCdcVdr;
}

void CdcCorrectionTask::SetCdcVdr(double value)
{
	fCdcVdr = value;
}




ClassImp(CdcCorrectionTask)