#include "TpcClusterFinderSimple.h"

#include <iostream>
#include <algorithm>
#include <stdexcept>
#include "assert.h"

#include "TpcPadPlane.h"
#include "TpcCluster.h"
#include "TpcDigiMapper.h"
#include "TpcDigiAmplitude.h"
#include "McId.h"
#include "McIdCollection.h"
#include "TpcPrelimCluster.h"

#ifdef TESTCHAMBER
#include "Pedestals.h"
#endif


bool sortPrelimClustersAmp(TpcPrelimCluster* lhs, TpcPrelimCluster* rhs){
	return (lhs->getAmp() > rhs->getAmp());
}

TpcClusterFinderSimple::TpcClusterFinderSimple(TpcPadPlane* p,
		std::vector<TpcCluster*>* ob,
		unsigned int timeslice, double G, double C)
: fpadplane(p), foutput_buffer(ob), fdt(timeslice), noXclust(false),
  splitDigis(0), fG(G), fC(C), maxClusterSlice(4000), sectorize(true), fnSectors(1), zJitter(0.1),
  fClusterMerge(true),frecalcDigiPos(kFALSE)
{
	// construct sector map
	std::vector<unsigned int> ids=fpadplane->GetSectorIds();
	fnSectors = ids.size();
	for(unsigned int is=0; is<fnSectors; ++is){
		unsigned int Sectorid = ids[is];
		fsectormap[Sectorid] = new std::vector<TpcDigi*>();
	}
	std::cout<<"TpcClusterFinderSimple: "
			<<fnSectors<<" Sectors instantiated."<<std::endl;

	if (fnSectors==1) return; // no need to build fpadIDsectorIDmap


	// build a map from padIds to SectorIds,
	// the purpose is to put digis at the borders of sectors into the digi buffers of both sectors
	// so that clusters won't be split at sector borders
	std::cerr<<"TpcClusterFinderSimple::TpcClusterFinderSimple - building fpadIDsectorIDmap ";

	unsigned int nPads, nNeighbors, sectorID;
	TpcPad* pad;
	std::vector<unsigned int> borderPads; // pads that lie at sector borders

	nPads = fpadplane->GetNPads();

	// loop over pads and find pads at borders
	for (unsigned int iPad=0; iPad<nPads; ++iPad){
		pad = fpadplane->GetPad(iPad);
		sectorID = pad->sectorId();

		fpadIDsectorIDmap[iPad] = new std::vector<unsigned int> (1, sectorID);

		//loop over neighbors
		nNeighbors = pad->nNeighbours();
		for (unsigned int ineigh=0; ineigh<nNeighbors; ++ineigh){
			unsigned int neighID = fpadplane->GetPad(pad->getNeighbour(ineigh))->sectorId();
			if (neighID == sectorID) continue;

			borderPads.push_back(iPad);
			break;
		}
	}

	double overlap(1.1);
	double x, y;
	unsigned int padId, padIdNeigh, nPadsInArea;;


	// loop over digis at borders
	for (unsigned int i=0; i<borderPads.size(); ++i){
		padId = borderPads[i];

		sectorID = fpadplane->GetPad(padId)->sectorId();

		std::vector<TpcPad*> buffer;
		fpadplane->GetPadXY(padId, x, y);
		fpadplane->GetPadList(x, y, overlap, buffer);

		// loop over pads in area
		nPadsInArea = buffer.size();
		for (unsigned int iPad=0; iPad<nPadsInArea; ++iPad){
			padIdNeigh = buffer[iPad]->padId();

			if ( (std::find(fpadIDsectorIDmap[padIdNeigh]->begin(), fpadIDsectorIDmap[padIdNeigh]->end(), sectorID)) == fpadIDsectorIDmap[padIdNeigh]->end() ){
				fpadIDsectorIDmap[padIdNeigh]->push_back(sectorID);
			}
		}
	}

	std::cout<<"... done"<<std::endl;


	//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;
	}
	zJitter = zDiff2.z() - zDiff1.z();
	//end of z jitter

}


TpcClusterFinderSimple::~TpcClusterFinderSimple(){
	std::map<unsigned int,std::vector<TpcDigi*>*>::iterator secIt=fsectormap.begin();
	while(secIt!=fsectormap.end()){
		delete secIt->second;
		++secIt;
	}
	fsectormap.clear();
}


void 
TpcClusterFinderSimple::process(std::vector<TpcDigi*>& alldigis)
{
	unsigned int nalldigi = alldigis.size();
	if(nalldigi<3) return;

	if(sectorize && fnSectors>1){
		// reserve
		std::map<unsigned int,std::vector<TpcDigi*>* >::iterator secIt=fsectormap.begin();
		while(secIt!=fsectormap.end()){ // loop over sectors
			secIt->second->reserve(nalldigi/fsectormap.size()+250);
			++secIt;
		} // end loop over sectors

		std::vector<unsigned int>* sectorIDs;

		// sectorize on pad plane
		for(int idi=0;idi<nalldigi;++idi){ // loop over digis
			// loop over sectorIDs
			sectorIDs = fpadIDsectorIDmap[alldigis[idi]->padId()];
			for(unsigned int i=0; i<sectorIDs->size(); ++i){
				fsectormap[(*sectorIDs)[i]]->push_back(alldigis[idi]);
			}
		} // end loop over digis

		// now process each sector independently
		secIt=fsectormap.begin();
		while(secIt!=fsectormap.end()){ // loop over sectors
			processSector(*(secIt->second), secIt->first);
			secIt->second->clear(); // clean up
			++secIt;
		} // end loop over sectors
	}
	else {
		processSector(alldigis, -1); // do not sectorize!
	}

}


void
TpcClusterFinderSimple::processSector(std::vector<TpcDigi*>& digis, int sectorID) {

	std::sort(digis.begin(), digis.end(), TpcDigiAge()); // for sectorizing in z

	unsigned int ndigi=digis.size();
	if(ndigi<3) return;

	// sectorize in z
	unsigned int nSlices = ndigi/maxClusterSlice + 1;
	unsigned int clusterSlice = ndigi/nSlices + 1;

	double startTime, stopTime, time;
	unsigned int firstDigi, lastDigi;

	for (unsigned int iSlice=0; iSlice<nSlices; ++iSlice){
		// find passive time window
		firstDigi = iSlice*clusterSlice;
		lastDigi = (iSlice+1)*clusterSlice - 1;
		if (lastDigi > ndigi-2) lastDigi = ndigi-1;

		startTime = digis[firstDigi]->t();
		stopTime = digis[lastDigi]->t();

		while(firstDigi>0){
			time = digis[--firstDigi]->t();
			if (fabs(startTime-time) > 2*fdt) break;
		}
		while(lastDigi<ndigi-1){
			time = digis[++lastDigi]->t();
			if (fabs(time-stopTime) > 2*fdt) break;
		}

		std::vector<TpcDigi*> digisInSlice;
		unsigned int ndigisInSlice = lastDigi-firstDigi;
		digisInSlice.reserve(ndigisInSlice);

		for (unsigned int i=firstDigi; i<lastDigi; ++i){
			digisInSlice.push_back(digis[i]);
		}


		std::stable_sort(digisInSlice.begin(),digisInSlice.end(),TpcDigiAmplitude());

		std::vector<TpcPrelimCluster*> prelimClusters;

		int prelimClusterCounter=0;

		prelimClusters.push_back(new TpcPrelimCluster(fpadplane, getTimeslice(digisInSlice.back()), prelimClusterCounter++, fG, fC));
		prelimClusters[0]->addHit(digisInSlice.back(), noXclust);

		// first pass
		for(unsigned int idigi = ndigisInSlice-2; ; --idigi) { // loop over digis from back to front, last digi was already processed
			std::vector<unsigned int> selClusters; // contains indices of clusters that the digi might belong to
			std::vector<unsigned int> selClustersShared; // contains indices of clusters that the digi might belong to
			TpcDigi* digi = digisInSlice[idigi];
			for(unsigned int iclust=0;iclust<prelimClusters.size();++iclust) { // loop over prelimClusters
				int inCluster = prelimClusters[iclust]->isInCluster(digi);
				if(inCluster == 0) {}
				else if (inCluster == 1)  selClusters.push_back(iclust);
				else selClustersShared.push_back(iclust);
			} // end loop over prelimClusters
			unsigned int nselclust = selClusters.size();

			if(nselclust == 0){
				unsigned int nselclustShared = selClustersShared.size();
				if (nselclustShared > 0){ // digi only neighbours to shared digis -> split between corresponding clusters
					double share = 1./(double)nselclustShared;
					for(unsigned int i=0; i<nselclustShared; ++i) {
						prelimClusters[selClustersShared[i]]->addHit(digi, noXclust, share);
						// add information which clusters are shared to other clusters
						for(unsigned int j=0; j<nselclustShared; ++j) {
							if (i!=j) prelimClusters[selClustersShared[j]]->addSharedCluster(digi, prelimClusters[selClustersShared[i]]);
						}
					}
					splitDigis += nselclustShared-1;
				}
				else { // digi cannot be added to existing cluster -> create a new cluster
					TpcPrelimCluster* fc = new TpcPrelimCluster(fpadplane, getTimeslice(digi), prelimClusterCounter++, fG, fC);
					fc->addHit(digi, noXclust);
					prelimClusters.push_back(fc);
				}
			}
			else if(nselclust == 1) { // digi can only belong to one cluster -> add to this cluster
				prelimClusters[selClusters[0]]->addHit(digi, noXclust);
			}
			else { // digi can belong to more than one cluster -> assign digi to these clusters with corresponding share value
				double share = 1./(double)nselclust;
				for(unsigned int i=0; i<nselclust; ++i) {
					prelimClusters[selClusters[i]]->addHit(digi, noXclust, share);
					// add information which clusters are shared to other clusters
					for(unsigned int j=0; j<nselclust; ++j) {
						if (i!=j) prelimClusters[selClusters[j]]->addSharedCluster(digi, prelimClusters[selClusters[i]]);
					}
				}
				splitDigis += nselclust-1;
			}

			if(idigi==0) break;
		} // end loop over digis

		// second pass: merge small clusters (1 unshared digi) into others
		if (fClusterMerge){
			std::sort(prelimClusters.begin(), prelimClusters.end(), sortPrelimClustersAmp); // sort by decreasing amp
			TpcPrelimCluster* cl;
			TpcPrelimCluster* clS;
			for(unsigned int iclust=prelimClusters.size()-1; iclust>0; --iclust) { // go from last to first because prelimClusters.erase will be faster
				cl = prelimClusters[iclust];
				if (cl->getNUnsharedDigis()>1) continue;
				unsigned int nDigis = cl->getNDigis();
				if (nDigis==1) continue;

				// add the digis to the other clusters
				std::set<TpcPrelimCluster*> sharedClusters;
				TpcDigi* d;
				for (unsigned int iDigi=nDigis-1; iDigi!=0; --iDigi){
					d = cl->getDigi(iDigi);
					unsigned int nShCl = cl->getNSharedClusters(d);
					double share = 1./double(nShCl);
					// put the digis to the clusters which share them
					for (unsigned int iSharedClust=0; iSharedClust<nShCl; ++iSharedClust) {
						clS = cl->getSharedCluster(d, iSharedClust);
						if (cl != clS) {
							clS->addHit(d, noXclust, share, true);
							sharedClusters.insert(clS);
						}
						else clS->removeSharedCluster(d, cl);
					}
				}

				// the last digi is the unshared one -> share it between all clusters
				d = cl->getDigi(0);
				double share = 1./double(sharedClusters.size());

				std::set<TpcPrelimCluster*>::iterator it;
				for (it=sharedClusters.begin(); it!=sharedClusters.end(); ++it){
					(*it)->addHit(d, noXclust, share, true);
				}

				// erase the cluster
				prelimClusters.erase(prelimClusters.begin()+iclust);
			} // end loop over prelimClusters
		} // end if (fClusterMerge)

		// convert prelimClusters to TpcClusters
		TpcDigi* digi0;
		for(unsigned int i=0;i<prelimClusters.size();++i)
		{
			digi0 = prelimClusters[i]->getDigi(0);
			double t = digi0->t();
			if ((iSlice < nSlices-1 && t>=startTime && t<stopTime) ||
					(iSlice == nSlices-1 && t>=startTime && t<=stopTime))
			{
				if (sectorID == -1)
				{ // no sectorization
					prelimClusters[i]->setRecalcDigiPos(frecalcDigiPos);
					prelimClusters[i]->setDoPRF(fDoPRF);
					foutput_buffer->push_back(prelimClusters[i]->convTpcCluster(0, zJitter, fsaveRaw)); // every cluster has secID 0
					//std::cout<<"clusterfinder clusterpos: "<<foutput_buffer->back()->pos().X()<<" "<<foutput_buffer->back()->pos().Y()<<" "<<foutput_buffer->back()->pos().Z()<<"\n";
				}
				else if (fpadplane->GetPad(digi0->padId())->sectorId() == sectorID)
				{
					prelimClusters[i]->setRecalcDigiPos(frecalcDigiPos);
					prelimClusters[i]->setDoPRF(fDoPRF);
					foutput_buffer->push_back(prelimClusters[i]->convTpcCluster(sectorID, zJitter, fsaveRaw));
				}
			}
			delete prelimClusters[i];
		}
		prelimClusters.clear();
	} // end loop over time-slices

}


double
TpcClusterFinderSimple::getTimeslice(TpcDigi* d){
	//double driftlen = TpcDigiMapper::getInstance()->z_from_tick(d->t()) - TpcDigiMapper::getInstance()->zGem();
	double driftlen = TpcDigiMapper::getInstance()->z_from_tick(d->t());
	double diffSigma = TpcDigiMapper::getInstance()->getGas()->Dl() * sqrt(driftlen);
	double nSamp = diffSigma/zJitter; // translate diffSigma into samples
	return fdt*nSamp + fdt; // + constant offset
}


void 
TpcClusterFinderSimple::reset()
{ 
	splitDigis = 0;
}


void
TpcClusterFinderSimple::checkConsistency()
{
	std::cout << "TpcClusterFinderSimple::checkConsistency() empty implementation" << std::endl;
}