#include "PndTpcClusterFinderSimple.h"

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

#include "PndTpcPadPlane.h"
#include "PndTpcCluster.h"
#include "PndTpcDigiMapper.h"
#include "PndTpcDigiAmplitude.h"
#include "McId.h"
#include "McIdCollection.h"

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


PndTpcPrelimCluster::PndTpcPrelimCluster(PndTpcPadPlane* p, double t, int id, double G, double C) :
  famp(0.), fcogT(-1.), fpadplane(p), ftimeslice(t), fid(id), fG(G), fC(C)
{
}


PndTpcPrelimCluster::~PndTpcPrelimCluster()
{
}


void PndTpcPrelimCluster::addHit(PndTpcDigi* digi, bool noXclust, double share) {
  fdigis.push_back(digi);
  fdigiShares[digi] = share;

  fcogT=0.;
  famp=0;
  for(unsigned int i=0;i<fdigis.size();++i) {
	  fcogT+=fdigis[i]->amp()*fdigiShares[fdigis[i]] * fdigis[i]->t();
	  famp+=fdigis[i]->amp()*fdigiShares[fdigis[i]];
  }
  fcogT*=1./famp;

  fpossiblePads.insert(digi->padId());	
  PndTpcPad* pad = fpadplane->GetPad(digi->padId());

  for(unsigned int ineigh=0; ineigh<pad->nNeighbours(); ++ineigh){
    PndTpcPad* neigh = fpadplane->GetPad( pad->getNeighbour(ineigh) );

    bool useThisNeighbor=true;
    if(noXclust){
      if(fabs(pad->x()-neigh->x())>0.01) useThisNeighbor=false;//gt 0.1mm
    }
    if(useThisNeighbor) fpossiblePads.insert(pad->getNeighbour(ineigh));	
  }

}


bool PndTpcPrelimCluster::isInTimeWindow(const PndTpcDigi* const digi){
  if( fabs(digi->t() - fcogT) <= ftimeslice ) 
    return true;
  return false;
}


bool PndTpcPrelimCluster::isInCluster(const PndTpcDigi* const digi) {
  if(!isInTimeWindow(digi))
    return false;
  int padID = digi->padId();
  if(fpossiblePads.count(padID) == 0) 
    return false;
  return true;
}


PndTpcCluster* PndTpcPrelimCluster::convPndTpcCluster(unsigned int sectorId, double zJitter, bool saveRaw) {
  cog(zJitter); // also error is calculated here
  PndTpcCluster* c = new PndTpcCluster(fpos,ferr,famp,fid,fdigis.size());
  c->SetMcId(fmcidCol);
  
  if(saveRaw){//defined in PndTpcAbsClusterFinder.h and default false
    for(unsigned int i=0;i<fdigis.size();++i){
      c->addDigi(fdigis[i], fdigiShares[fdigis[i]]);
    }
  }
  
  c->SetSector(sectorId);
  
  return c;
}


void PndTpcPrelimCluster::cog(double zJitter){
  bool DEBUG=false;

  fpos.SetXYZ(0,0,0);
  famp=0;
  McIdCollection mcid;
  unsigned int ndigis=fdigis.size();

  // loop over digis to calculate cog
  for(unsigned int id=0;id<ndigis;++id){
	  PndTpcDigi* adigi=fdigis[id];
	  double a=(double)adigi->amp()*fdigiShares[adigi];
	  mcid.AddIDCollection(adigi->mcId(),a);
	  TVector3 thispos;
	  PndTpcDigiMapper::getInstance()->map(adigi,thispos);
	  fpos+=a*thispos;
	  famp+=a;
  }
  fpos*=1./famp;


  // calculate errors: ------------------------------------------
	double dx, dy;
	PndTpcDigiMapper::getInstance()->padsize(fdigis[0]->padId(),dx,dy);

  ferr=(0.,0.,0.);
     
  for(unsigned int id=0;id<ndigis;++id){
    PndTpcDigi* adigi=fdigis[id];
    double a=(double)adigi->amp()*fdigiShares[adigi];
    TVector3 thispos;
    PndTpcDigiMapper::getInstance()->map(adigi,thispos);
    TVector3 df=thispos-fpos;
    double sigmaX_sq = a*df.X()*df.X();
    double sigmaY_sq = a*df.Y()*df.Y();
    double sigmaZ_sq = a*df.Z()*df.Z();
    
    TVector3 thissig(sigmaX_sq,sigmaY_sq,sigmaZ_sq);
    ferr+=thissig;  
  } // end second loop over digis

  if(ferr.X()<1E-5) ferr.SetX(sqrt(dx*dx/12)*fC/famp);
  else ferr.SetX(sqrt(ferr.X()/famp)*fC/famp);

  if(ferr.Y()<1E-5) ferr.SetY(sqrt(dy*dy/12)*fC/famp);
  else ferr.SetY(sqrt(ferr.Y()/famp)*fC/famp);

  if(ferr.Z()<1E-5) ferr.SetZ(sqrt(zJitter*zJitter/12)*fC/famp);
  else ferr.SetZ(sqrt(ferr.Z()/famp)*fC/famp);
  
  if(DEBUG) ferr.Print();

  fmcidCol = mcid;
}



PndTpcClusterFinderSimple::PndTpcClusterFinderSimple(PndTpcPadPlane* p,
						     std::vector<PndTpcCluster*>* 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)
{
  // 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<PndTpcDigi*>();
  }
  std::cout<<"PndTpcClusterFinderSimple: "
	   <<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<<"PndTpcClusterFinderSimple::PndTpcClusterFinderSimple - building fpadIDsectorIDmap ";

  unsigned int nPads, nNeighbors, sectorID;
  PndTpcPad* 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<PndTpcPad*> 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;

  // print
  /*std::vector<unsigned int>* sectorIDs;
  for (unsigned int iPad=0; iPad<nPads; ++iPad){
    std::cout<<iPad<<"  ";
    sectorIDs = fpadIDsectorIDmap[iPad];
    for(unsigned int i=0; i<sectorIDs->size(); ++i){
      std::cout<<(*sectorIDs)[i]<<" ";
    }
    std::cout<<"\n";
  }*/


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

  TVector3 zDiff1,zDiff2;
  PndTpcDigi zDiffDigi1(1,1,1,dummyColl),zDiffDigi2(1,2,1,dummyColl);
  PndTpcDigiMapper::getInstance()->map(&zDiffDigi1,zDiff1);
  PndTpcDigiMapper::getInstance()->map(&zDiffDigi2,zDiff2);
  zJitter = zDiff2.z() - zDiff1.z();
  //end of z jitter

}


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


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

  if(sectorize && fnSectors>1){
    std::cerr<<"processing digis sectorwise"<<std::endl;
    // reserve
    std::map<unsigned int,std::vector<PndTpcDigi*>* >::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 {
    std::cerr<<"processing all digis"<<std::endl;
    processSector(alldigis, -1); // do not sectorize!
  }

}


void
PndTpcClusterFinderSimple::processSector(std::vector<PndTpcDigi*>& digis, int sectorID)
{
  std::sort(digis.begin(),digis.end(),PndTpcDigiAge()); // 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();

    //std::cerr<<" active volume from "<<firstDigi<<" to "<<lastDigi<<std::endl;
    
    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::cerr<<"  passive volume from "<<firstDigi<<" to "<<lastDigi<<std::endl;

    std::vector<PndTpcDigi*> 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(),PndTpcDigiAmplitude());

    std::vector<PndTpcPrelimCluster*> prelimClusters;

    int prelimClusterCounter=0;

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

    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
      for(unsigned int iclust=0;iclust<prelimClusters.size();++iclust) { // loop over prelimClusters
        if(prelimClusters[iclust]->isInCluster(digisInSlice[idigi]) ) {
          selClusters.push_back(iclust);
        }
      } // end loop over prelimClusters
      unsigned int nselclust = selClusters.size();

      if(nselclust == 0){ // digi cannot be added to existing cluster -> create a new cluster
        PndTpcPrelimCluster* fc = new PndTpcPrelimCluster(fpadplane,fdt,prelimClusterCounter++, fG,fC);
        fc->addHit(digisInSlice[idigi],noXclust);
        prelimClusters.push_back(fc);
      }
      else if(nselclust == 1) { // digi can only belong to one cluster -> add to this cluster
        prelimClusters[selClusters[0]]->addHit(digisInSlice[idigi],noXclust);
      }
      else { // digi can belong to more than one cluster -> assign digi to these clusters with corresponding share value
        PndTpcDigi* fd = digisInSlice[idigi];
        double share = 1./(double)nselclust;
        for(unsigned int i=0; i<nselclust; ++i) {
          prelimClusters[selClusters[i]]->addHit(fd,noXclust, share);
        }
        splitDigis += nselclust-1;
      }

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

    // convert prelimClusters to PndTpcClusters
    PndTpcDigi* digi0;
    for(unsigned int i=0;i<prelimClusters.size();++i){
      digi0 = prelimClusters[i]->getDigi(0);
      double t = digi0->t();
      if (t>=startTime && t<stopTime){
        if (sectorID == -1){ // no sectorization
          foutput_buffer->push_back(prelimClusters[i]->convPndTpcCluster(0, zJitter, fsaveRaw)); // every cluster has secID 0
        }
        else if (fpadplane->GetPad(digi0->padId())->sectorId() == sectorID){
          foutput_buffer->push_back(prelimClusters[i]->convPndTpcCluster(sectorID, zJitter, fsaveRaw));
        }
      }
      delete prelimClusters[i];
    }
    prelimClusters.clear();
  } // end loop over time-slices

}


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


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