//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcdEdxTask
//      see TpcdEdxTask.hh for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Felix Boehmer, Christian Hoeppner     TUM         
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "TpcdEdxTask.h"

// C/C++ Headers ----------------------
#include <cmath>
#include <iostream>
#include <map>
#include <assert.h> 
#include <stdexcept>

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "TpcCluster.h"
#include "TpcDigiMapper.h"
#include "TpcFrontend.h"
#include "TpcSPHit.h"
#include "PndMCTrack.h"
#include "TpcPoint.h"
#include "TpcAlignmentManager.h"

#include "TpcPadShapePool.h"
#include "TpcPadPlane.h"
#include "TpcAbsPadShape.h"
#include "TpcPad.h"

#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "TpcDigiPar.h"


#include "GFAbsTrackRep.h"
#include "GFRecoHitFactory.h"
#include "GFKalman.h"
#include "GFException.h"
#include "TH1D.h"
#include "TFile.h"
#include "TGeoTrack.h"

#include "GFDetPlane.h"
#include "TpcdEdx.h"

#include "RecoHits/GFAbsRecoHit.h"
#include "TVector3.h"

//#include <signal.h>
//#include <stdlib.h>

#include "PndDetectorList.h"

// Class Member definitions -----------


TpcdEdxTask::TpcdEdxTask()
  : FairTask("dE/dx Task"), _persistence(kFALSE),
    _DX(1.), _idealdEdx(kFALSE), fUseFirstDigiPos(kFALSE),
    fIgnoreEdgeDigis(kFALSE)
{
  _trackBranchName = "TpcTrackPostFit";
  _clusterBranchName = "TpcCluster";
  _digiBranchName = "TpcDigi";
  _outBranchName = "TpcdEdx";
  //_spHitBranchName = "TpcSPHit";
  fVerbose = 0;
  //fPadPlaneFile = "tpc/FOPI/par/padplane_FOPI.dat";
  //fShapeFile = "tpc/parfiles/PadShapePrototype.dat";
}


TpcdEdxTask::~TpcdEdxTask()
{
}

InitStatus
TpcdEdxTask::Init()
{
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();
  
  if(ioman==0)
    {
      Error("TpcdEdxTask::Init","RootManager not instantiated!");
      return kERROR;
    }
  
  // Get input collection
  _trackArray=(TClonesArray*) ioman->GetObject(_trackBranchName);
  if(_trackArray==0)
    {
      Error("TpcdEdxTask::Init","track-array not found!");
      return kERROR;
    }
  
  if(_idealdEdx) {
    _mcTrackArray=(TClonesArray*) ioman->GetObject("MCTrack");
    if(_mcTrackArray==0)
      Fatal("TpcdEdxTask::Init","MCTrack-array not found! Cannot calculate MC dEdx");
  
    _pointArray=(TClonesArray*) ioman->GetObject("TpcPoint");
    if(_pointArray==0)
      Fatal("TpcdEdxTask::Init","Point-array not found! Cannot calculate MC dEdx");
  }

  _clusterArray=(TClonesArray*) ioman->GetObject(_clusterBranchName);
  if(_clusterArray==0)
    {
      Error("TpcdEdxTask::Init","Cluster/Hit - array not found!");
      return kERROR;
    }
 
  _digiArray=(TClonesArray*) ioman->GetObject(_digiBranchName);
  if(_digiArray==0)
    {
      Error("TpcdEdxTask::Init","Digi-array not found!");
      return kERROR;
    }
  
  _dEdxOutArray = new TClonesArray(_outBranchName);
  ioman->Register("dEdx","Tpc",_dEdxOutArray,_persistence);
  
  _dEdxMCOutArray = new TClonesArray("TpcdEdx");
  ioman->Register("dEdx_MC","Tpc",_dEdxMCOutArray,_persistence);
  
  // setup histograms
  _distHist=new TH1D("distHist","Stepping lentghs",500,-3.,3.);
  _dirHist=new TH1D("dirHist","direction abundances",10,-2.,2.);

 _hitBranchID = FairRootManager::Instance()->GetBranchId(_clusterBranchName);

  std::cout<<"Initializing PadPlane ... "<<std::endl;
  //TpcPadShapePool* pool = fPar->getPadShapes();
  fPlane = fPar->getPadPlane();
  fRMin = fPar->getRMin();
  fRMax = fPar->getRMax();

  if(fIgnoreEdgeDigis) {
    fRMin = fRCutoff1;
    fRMax = fRCutoff2;
  }
 
  return kSUCCESS;
}


void
TpcdEdxTask::SetParContainers() {

  std::cout<<"TpcClusterFinderTask::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");

}



void
TpcdEdxTask::Exec(Option_t* opt)
{
  std::cout<<"TpcdEdxTask::Exec()"<<std::endl;
  FairRootManager* ioman= FairRootManager::Instance();

  static unsigned int nTotalDigis(0);
  static unsigned int nDoubleDigis(0);
  int evts(0);   

  _dEdxOutArray->Delete();
  _dEdxMCOutArray->Delete();
  

  Int_t ntracks=_trackArray->GetEntriesFast();
  
  if(ntracks>20000){
    std::cout<<"ntracks="<<ntracks<<" Evil Event! skipping"<<std::endl;
    return;
  }
  
  for(unsigned int itr=0;itr<ntracks;++itr){
    double totDist = 0.;
    if (fVerbose) std::cout<<"TpcdEdxTask::Exec(): starting track "<<itr<<std::endl;
    GFTrack* trk=(GFTrack*)_trackArray->At(itr);
    if (fVerbose) std::cout<<"*** Number of hits in track: "<<trk->getNumHits()<<" ***"<<std::endl;

    TpcdEdx dedx;
    TpcdEdx dedx_MC;
    
    GFAbsTrackRep* theRep = trk->getCardinalRep();
    if (theRep->getStatusFlag() != 0) {
      // write empty dEdx and continue with next trk
      new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);

      if(_idealdEdx) {
        new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
      }
      continue;
    }

    
    GFTrackCand cand = trk->getCand();
        
    bool useSPHits = false;
    if(TString(_clusterArray->GetClass()->GetName()) == "TpcSPHit") 
      useSPHits = true;
    
    //these are SORTED, first entry corresponds to track's start-hit
    std::vector<int> hitIDs = cand.getHitIDs(_hitBranchID);
        
    //GET MC INFORMATION FOR CROSS-CHECK
    //only works with IdealTracking and pure trackIDs
    std::vector<TpcPoint*> pointlist;
    if(_idealdEdx) {
      if (fVerbose) std::cerr<<"collect MC points"<<std::endl;
      //loop over MC points and collect
      for(int p=0; p<_pointArray->GetEntriesFast(); ++p) {
        int id = ((TpcPoint*)_pointArray->At(p))->GetTrackID();
        if(id==0) //only primary tracks
          pointlist.push_back((TpcPoint*)_pointArray->At(p));
      }
    }
    
    TVector3 pos,mom;
    TVectorD protState = theRep->getState();
    if (fVerbose) std::cerr<<"get pos and mom"<<std::endl;
    try{
      pos = theRep->getPos();
      mom = theRep->getMom();
    }
    catch(GFException& e){
      e.what();
      // write empty dEdx and continue with next trk
      new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);
      if(_idealdEdx) {
        new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
      }
      continue;
      //TODO: exception handling
    }
    
    TVector3 startDir = mom;
    startDir.SetMag(1.);
    
    TVector3 altFirstPos;
    
    //reference plane:
    GFDetPlane pl;
    pl.setO(pos);
    pl.setNormal(startDir);

    if (fVerbose) std::cerr<<"collect digis"<<std::endl;
    // get all digis in the track
    std::map<unsigned int, TpcDigi*> digis;   //<TCA index, pointer>
    unsigned int count = 0;
    TClonesArray* actClArr;
    int foundDigi = -1; //first digi in (opposite) track direction

    for(int i=0; i<hitIDs.size(); i++){ //loop over hits (clusters)
      TpcCluster* cl;
      if(useSPHits) {
	TpcSPHit* theHit = (TpcSPHit*) _clusterArray->At(hitIDs[i]);
	int clBrID = theHit->getClusterBranchID();
	int clID = theHit->getClusterID(); 
	if(count==0) {
	  TString clBrName = ioman->GetBranchName(clBrID);
	  actClArr = (TClonesArray*) ioman->GetObject(clBrName);
	}
	
	cl = (TpcCluster*)(actClArr->At(clID));
	count++;
      }
      else 
	cl = (TpcCluster*) _clusterArray->At(hitIDs[i]);
            
      const std::map<unsigned int,double>* digiMap = cl->getDigiMap();
            
      double maxDist = 0.;  //(cm)
     
      for(std::map<unsigned int, double>::const_iterator it=digiMap->begin();it!=digiMap->end();++it){
	//avoid double-counting
	if(!digis.count(it->first)) 
	  digis[it->first] = (TpcDigi*) _digiArray->At(it->first);
	else 
	  continue;
	TpcDigi* aDigi = digis[it->first];
	
	//only for first hit (cluster): find first Digi
	if(i==0 && fUseFirstDigiPos) {
	  //get Digi coordinates:
	  TVector3 dPos;
	  try {
	    TpcDigiMapper::getInstance()->map(aDigi,dPos);
	  }
	  catch (...) {
	    std::cerr<<"TpcdEdxTask::Could not get Position of digi! Skipping ..."<<std::endl;
	    continue; 
	  }
	  dPos = TpcAlignmentManager::getInstance()->localToMaster("tpc",dPos);
	  TVector3 pc, dpc;
    try{
      theRep->extrapolateToPoint(dPos,pc,dpc);
    }catch(GFException exp){
      std::cerr<<"TpcdEdxTask::Exec() error during first digi extrapolation:"<<std::endl;
      std::cerr<<exp.what()<<std::endl;
      continue;
    }
	  TVector3 dDist = pl.getO() - pc;
	  if(dDist.Mag() > maxDist && (dDist*pl.getNormal()) < 0.) {  //we want negative extrap!
	    if(fIgnoreEdgeDigis)
	      if(dPos.Perp() < fRCutoff1 || dPos.Perp() > fRCutoff2)
		continue;
	    foundDigi = (*it).first;
	    altFirstPos = dPos;
	    maxDist = dDist.Mag();
	  }
	}
      } //end loop over this cluster's digis
   
      
    } //end loop over hits (clusters)
  
    
    //startDir *= -1.;
    dedx.setMom(mom);
    
    if (fVerbose) {
      TVector3 pos0(((TpcCluster*)(_clusterArray->At(hitIDs[0])))->pos());
      std::cout<<"pos - pos0 mag"<< (pos-pos0).Mag() << std::endl;
    }
    
    GFDetPlane here,next;
    here.setO(pos);
    here.setNormal(mom);
    if(fVerbose) {std::cout<<"Original extrap direction:"<<std::endl;mom.Print();std::cout<<std::endl;}
    //if we successfully found a "first digi" use it as starting pos
    if (foundDigi >= 0 && fUseFirstDigiPos) {
      if(fVerbose) {
	std::cout<<"Moved plane from cluster position (r="<<here.getO().Perp()<<"): \n    " ;
	here.getO().Print();
	std::cout.flush();
	std::cout<<"\nto first Digi position (r="<<altFirstPos.Perp()<<"): \n    " ;	
	altFirstPos.Print();
	std::cout.flush();
	std::cout<<"\n"<<std::endl;
      }      
      double pull = 0.99;
      if(altFirstPos.Perp() > 2*fRMin)  //first digi outside, track supposedly goes inwards
	pull = 1.01; 
      here.setO(altFirstPos*pull);  //pull it inward a little
    }
      
    
    int counter(0);    
    bool exc = false;
        
    std::set<unsigned int> usedDigis;  //bookkeeping for speedup
    
    while(true){ //stepalong until no digi is found between the planes
      bool abort = true;
      bool forceabort = false;
      double dir = 1.;
      if (fVerbose) std::cerr<<"at counter "<<counter<<std::endl;
      if(counter++>0) here=next;//not the first time
      else if(here.getNormal() * startDir < 0) {
        //_DX*=-1;
        dir=-1.;
      }
      TVector3 destination = here.getO()+(here.getNormal()*_DX*dir);
            
      TVector3 poca,dirInPoca;
      double dist;
      try{
        theRep->extrapolateToPoint(destination,poca,dirInPoca);
	next.setO(poca);
	next.setNormal(dirInPoca);
      }
      catch(GFException& e){
        e.what();
        exc=true;
        break;
        //TODO: exception handling
      }

      //check if next lies out of the TPC volume
      
      if(next.getO().Perp() > fRMax || next.getO().Perp() < fRMin) {
	if(fVerbose)
	  std::cout<<"OUT OF CHAMBER  next.Perp(): <<"<<next.getO().Perp()<<std::endl;
	TVector3 conn = next.getO() - here.getO();
	conn.SetMag(1.);
	double drn, drh, dr;
	dr = next.getO().Perp() - here.getO().Perp();
	double comp = fRMax;
	if(next.getO().Perp()<fRMin)
	  comp = fRMin;
	drn = next.getO().Perp() - comp;
	double moveR = 1.1*fabs(drn)*_DX/fabs(dr);   //scaled to get safely inside the cage
	TVector3 newpos = next.getO() - (conn*moveR);
	next.setO(newpos);
	if(fVerbose)
	  std::cout<<"AFTER SHIFT  next.Perp(): <<"<<next.getO().Perp()<<std::endl;
	forceabort = true;
      }

      //now do the final dE/dx extrapolation
      try {
	dist = theRep->extrapolate(next,protState);
      }
      catch(GFException& e) {
        e.what();
        exc=true;
        break;
      }
      
      double thisDist = dist-totDist;
      totDist = dist;
      
      if (fVerbose) {
	std::cout<<"Extrapolation results :\n";poca.Print();
	std::cout<<std::endl;dirInPoca.Print();
	std::cout<<"Extrap distance: "<<dist<<std::endl;
      }
      
      TVector3 normHere = here.getNormal();
      TVector3 normNext = next.getNormal();
          
      double dE = 0.;
      
      double behindHere;
      double behindNext;
    

      if (fVerbose) {
	here.Print();
	next.Print();
      }

      std::map<unsigned int, TpcDigi*>::const_iterator cit;
      for(cit=digis.begin(); cit!=digis.end();cit++){
        if(usedDigis.count(cit->first))
	   continue;
        TVector3 digiPos;
	try {
	  TpcDigiMapper::getInstance()->map(cit->second,digiPos);
	}
	catch(const std::exception& ex) {
	  std::cout<<ex.what()<<std::endl;
	  continue; //THIS WAS MISSING!
	}
	if(fIgnoreEdgeDigis)
	  if(digiPos.Perp() < fRCutoff1 || digiPos.Perp() > fRCutoff2)
	    continue;
        digiPos=TpcAlignmentManager::getInstance()->localToMaster("tpc",digiPos);

        behindHere = normHere * (here.dist(digiPos));
        behindNext = normNext * (next.dist(digiPos));
	
	behindHere *= 1./fabs(behindHere);
	behindNext *= 1./fabs(behindNext);

	
	double ok = behindHere + behindNext;    // 0 if digi between planes
	if (fVerbose) std::cout<<"JDIOSFSDF   OK:  "<<ok<<std::endl;
	if (behindHere < 0.) {
	   abort = false;  //there are still viable digis up fronts, keep running
	   if(fabs(ok) < 0.5) {  //possibilities are actually -2, 0 ,2 within double precision...
	     dE+=(cit->second)->amp();
	     usedDigis.insert(cit->first);
	   }
	}
      }

      if (fVerbose) std::cout<<"dE: "<<dE<<",  dist: "<<dist<<",  abort: "<<abort<<std::endl;
      if(abort || forceabort) 
	break;
      
      if(dE >= 0.)
	dedx.add(dE,thisDist);

      //MC PART ---------------------------------------------------
      
      dE = 0.;
      
      if(_idealdEdx) {
        for(unsigned int m=0;m<pointlist.size();++m){
          TVector3 pointPos;
          pointlist.at(m)->Position(pointPos);
          if((here.getO()-pointPos).Mag() > (_DX+2.)) continue;
          behindHere = normHere * (here.dist(pointPos));
          behindNext = normNext * (next.dist(pointPos));
          if(behindHere<0.){//in front of
            if(behindNext>=0.){
              dE+=pointlist.at(m)->GetEnergyLoss();
            }
          }
          else{
            if(behindNext<0.){//in front of
      	      dE+=pointlist.at(m)->GetEnergyLoss();
            }
          }
        }
	
        if(dE>0. && (!exc)){
          dedx_MC.add(dE*1E9,dist);
        }
      } // end if _idealdEdx
    }
    
    if(exc==false) 
      new((*_dEdxOutArray)[itr]) TpcdEdx(dedx);
    else
      new((*_dEdxOutArray)[itr]) TpcdEdx();

    if(_idealdEdx) {
      new((*_dEdxMCOutArray)[itr]) TpcdEdx(dedx_MC);
    }
    //for(unsigned int i=0;i<digis.size();++i){
    //  delete digis[i];
    //}
    
  }
}


void 
TpcdEdxTask::WriteHistograms(const TString& filename){
  TFile* file = new TFile(filename,"RECREATE");
  file->mkdir("DEDX");
  file->cd("DEDX");

  _distHist->Write();
  delete _distHist;

  _dirHist->Write();
  delete _dirHist;
   
  file->Close();
  delete file;
}

void 
TpcdEdxTask::SetIgnoreEdgeDigis(double r1, double r2, Bool_t opt) {
  fRCutoff1 = r1;
  fRCutoff2 = r2;
  fIgnoreEdgeDigis = opt;
}



ClassImp(TpcdEdxTask)