/////////////////////////////////////////////////////////////
//
//  FastSim
//
//  Reader for McHit
////
/////////////////////////////////////////////////////////////// 

//C++ class headers
#include <string>
#include <fstream>

//CBM class headers
#include "CbmRootManager.h"
#include "CbmRunAna.h"
#include "CbmRuntimeDb.h"
#include "CbmMCTrack.h"
#include "CbmStack.h"
#include "CbmRun.h"

#include "PndFastSim.h"

//Rho includes
#include "TCandidate.h"
#include "TSimpleVertex.h"

//ROOT class headers
#include "TClonesArray.h"
#include "TVector3.h"
#include "TParticle.h"
#include "TRandom3.h"
#include "TDatabasePDG.h"
#include "TParticlePDG.h"
#include "TVirtualMC.h"
#include "TF1.h"

//Fast Sim class headers
#include "PndFsmTrack.h"
#include "PndFsmResponse.h"
#include "PndFsmAbsDet.h"
#include "PndFsmDetFactory.h"

using std::cout;
using std::endl;
using std::string;
using std::ifstream;

		
// -----   Default constructor   -------------------------------------------
PndFastSim::PndFastSim() :
  CbmTask("FastSim Dump") { 
  //  fCandidates = new TClonesArray("TParticle");
  fRand=new TRandom3();
  fDetFac = new PndFsmDetFactory;
  fAddedDets=" ";
  fVb=0;
  fGenSplitOffs=false;
}
// -------------------------------------------------------------------------

// -----   Destructor   ----------------------------------------------------
PndFastSim::~PndFastSim() { 
    CbmRootManager *fManager =CbmRootManager::Instance();
   fManager->Write();

  if (fCandidates) {fCandidates->Delete(); delete fCandidates;}
  if (fMcCandidates) {fMcCandidates->Delete(); delete fMcCandidates;}

  if (fDetFac) delete fDetFac;
  if (fRand) delete fRand;
}
// -------------------------------------------------------------------------


void PndFastSim::Register() {
  //---

  fCandidates = new TClonesArray("TParticle");
  CbmRootManager::Instance()->Register("PndParticles","FastSim", fCandidates, kTRUE);

  fMcCandidates = new TClonesArray("TParticle");
  CbmRootManager::Instance()->Register("PndMcParticles","FastSim", fMcCandidates, kTRUE);

  fFsmCandidates = new TClonesArray("TCandidate");
  CbmRootManager::Instance()->Register("PndCandidates","FastSim", fFsmCandidates, kTRUE);

}



// -----   Public method Init   --------------------------------------------
InitStatus PndFastSim::Init() {
 
  if (fVb>3) cout << " Inside the Init function****" << endl;
  
  //CbmDetector::Initialize();
  //CbmRun* sim = CbmRun::Instance();
  //CbmRuntimeDb* rtdb=sim->GetRuntimeDb();
  
  // Get RootManager
  /* CbmRootManager* ioman = CbmRootManager::Instance();
  if ( ! ioman ) {
    cout << "-E- PndFastSim::Init: "
	 << "RootManager not instantiated!" << endl;
    return kFATAL;
  }
  */
  Register();
  if (fVb) 
  {
    cout <<"\n\n******   DETECTOR SETUP SUMMARY    ***************\n\n";
    for (FsmAbsDetList::iterator iter=fDetList.begin();iter!=fDetList.end(); iter++)
    {
      cout<<endl<<"----------------------------"<<endl;
      PndFsmAbsDet *det=*iter;
      det->print(std::cout);
    }    
    cout <<"\n\n******   DETECTOR SETUP SUMMARY    ***************\n\n";
  }
    
  // Create and register output array
  cout << "-I- PndFastSim: Intialization successfull" << endl;
  return kSUCCESS;

}

void PndFastSim::SetParContainers() {

  // Get run and runtime database
  CbmRunAna* run = CbmRunAna::Instance();
  if ( ! run ) Fatal("SetParContainers", "No analysis run");

  CbmRuntimeDb* db = run->GetRuntimeDb();
  if ( ! db ) Fatal("SetParContainers", "No runtime database");

 
}
// -------------------------------------------------------------------------

// -----   Enable split off parametrization   ------------------------------

bool PndFastSim::EnableSplitoffs(std::string fname)
{
  if (fname=="")
  {
    cout <<" -W-  (PndFastSim::EnableSplitoffs) - no filename given; no split offs will be produced."<<endl;
    return false;
  }

  
  /*
  prepare the split off parametrization
    expected in the parameter file:
    - four blocks (mom, phi, tht, multiplicity) with the parametrization structured like
    * <n_i>   <lowlimit>   <uplimit>   <model-string(in ROOT TF1 format)>
    * 5 lines of parameter sets with info for pid=11,13,211,321,2212
      <pid> <p1> <p2> ... <p_n_i>

    for the moment momentum independent
    models:
    momentum     : expo(0)+expo(2) = 4 params 
    tht+phi      : gaus(0)+gaus(3) = 6 params
    multiplicity : expo(0)+gaus(2) = 5 params
  */

  ifstream pars(fname.data());
  
  if ( (pars.rdstate() & ifstream::failbit ) != 0 )
  {
    cout << " -W-  (PndFastSim::EnableSplitoffs) -  Error opening '"<<fname<<"'; no split offs will be produced."<<endl;
    return false;
  }

  cout <<"tuudel"<<endl;
  
  int i,j,k;
  char tmp[50];

  string model[4];
  int numpars[4];
  double rangeup[4],rangelow[4];
  double curpar;
  
  
  //loop over blocks
  for (i=0;i<4;i++)
  { 
    pars>>model[i]>>numpars[i]>>rangelow[i]>>rangeup[i];
    if (fVb) cout <<" -I-  (PndFastSim::EnableSplitoffs)  split off model #"<<i<<": '"<<model[i]<<"' with "<<numpars[i]<<" params."<<endl;  
    
    //loop over pid
    for (j=0;j<5;j++)
    {
     sprintf(tmp,"f%d%d",j,i);
     fspo[j][i]=new TF1(tmp,model[i].data());
     
     //loop over params
     pars>>curpar;   // read the dummy PID value
     for (k=0;k<numpars[i];k++) 
     {
       pars>>curpar;
       fspo[j][i]->SetParameter(k,curpar);
     }
     
     fspo[j][i]->SetRange(rangelow[i],rangeup[i]);
     if (fVb) fspo[j][i]->Print();
    }
  }

  if (fVb) cout <<" -I-  (PndFastSim::EnableSplitoffs) - Successfully read "<<fname<<endl;  
  
  pars.close();

  fGenSplitOffs=true;
 
  return true;
}


// -------------------------------------------------------------------------

// -----   Add a detector to the setup   -----------------------------------
bool PndFastSim::AddDetector(std::string name, std::string params)
{
  if (name =="")
  {
    cout <<" -W-  (PndFastSim::AddDetector) - No name given, no detector added!"<<endl;

  }

  if (fAddedDets.find(" "+name+" ")!=std::string::npos) 
  {
    std::cout <<" -W-  (PndFastSim::AddDetector) - Detector <"<<name<<"> already appended - skipping!"<<std::endl;
    return false;
  }
  fAddedDets.append(std::string(name+" "));
  
  PndFsmAbsDet *det=0;
  det=fDetFac->create(name,params);
  if (!det) return false;
  
  fDetList.push_back(det);
  std::cout<<" -I- (PndFastSim::AddDetector) - Added detector "<<name<<" with params <"<<params<<">"<<std::endl;
  return true;
}

// -------------------------------------------------------------------------

// -----   Public method Exec   --------------------------------------------
void PndFastSim::Exec(Option_t* opt) {
 
  CbmStack *fStack=(CbmStack*)gMC->GetStack();  
  int nTracks=fStack->GetNtrack();


   // Reset output array
  if (fCandidates->GetEntriesFast() != 0)  fCandidates->Clear("C");
  if (fMcCandidates->GetEntriesFast() != 0)  fMcCandidates->Clear("C");
  if (fFsmCandidates->GetEntriesFast() != 0)  fFsmCandidates->Clear("C");

  TClonesArray &tracks = *fCandidates;
  TClonesArray &mctracks = *fMcCandidates;
  TClonesArray &candidates= *fFsmCandidates;

  if (fVb)cout <<"number of tracks **** "<< nTracks <<endl;
  
  for (Int_t iPoint=0; iPoint<nTracks; iPoint++){

	  // Get the MCTrack information
    Int_t size = tracks.GetEntriesFast();
    Int_t mcsize = mctracks.GetEntriesFast();
    Int_t candsize = candidates.GetEntriesFast();

    TParticle *t = fStack->GetParticle(iPoint);
    if (fVb>1) t->Print();
    
    TLorentzVector p4(t->Px(),t->Py(),t->Pz(),t->Energy());

    TVector3 stvtx(t->Vx(),t->Vy(),t->Vz());

    TLorentzVector vtx(stvtx,t->T());
    
    double charge=TDatabasePDG::Instance()->GetParticle(t->GetPdgCode())->Charge()/3.;

    PndFsmTrack *ft=new PndFsmTrack(p4,stvtx,stvtx,charge,t->GetPdgCode(),iPoint+1);

    TCandidate *tcand=new (candidates[candsize]) TCandidate(p4,charge,new TSimpleVertex(stvtx));


    // store a plain copy of the mc track to the file
    TParticle *pmc=new (mctracks[mcsize]) TParticle(ft->pdt(),0,0,0,0,0,ft->p4(),TLorentzVector(ft->startVtx(),t->T()));
    
    //    PndFsmTrack *tt=new (fsmtracks[mcsize]) PndFsmTrack(p4,stvtx,stvtx,charge,t->GetPdgCode(),iPoint+1);

    // smear and cut the track according to the detector setup
    if (smearTrack(ft))
    {
      TParticle *p=new (tracks[size]) TParticle(ft->pdt(),0,0,0,0,0,ft->p4(),TLorentzVector(ft->startVtx(),t->T()));
      TCandidate *tcand=new (candidates[candsize]) TCandidate(ft->p4(),ft->charge(),new TSimpleVertex(ft->startVtx()));
    }
     // TParticle(mctra->GetPdgCode(),0,0,0,0,0,p4,vtx);

    /*    cout <<"pdg("<<mctr->GetPdgCode()<<") ";
    cout <<"mothid("<<mctr->GeCode()<<") ";
   cout <<"p3("<<mctr->GetMomentum().Mag()<<") ";
    cout<<"Momentum "<<mctra->GetMomentum().X()<<" "<<mctra->Get4Momentum().X()<<endl;
    */
  }
}
// -------------------------------------------------------------------------

bool PndFastSim::smearTrack(PndFsmTrack *t)
{

  
  FsmResponseList responseList;
  
  for (FsmAbsDetList::iterator iter=fDetList.begin();iter!=fDetList.end(); iter++)
  {
    PndFsmAbsDet *det=*iter;
    //if (!det) {cout <<"--------------------------> outch"<<endl;continue;}

    PndFsmResponse *respo=det->respond(t);
    
    if (respo) responseList.push_back(respo);
  }   
    
  PndFsmResponse *allResponse=sumResponse(responseList);

  t->setDetResponse(allResponse);

  for (FsmResponseList::iterator riter=responseList.begin(); riter!=responseList.end();riter++) 
  {
    PndFsmResponse *resp=*riter;
    if (resp) delete resp;
  }    
  responseList.clear();
 

  return cutAndSmear(t);

  /*
  TLorentzVector l=t.p4();
    
    l.SetPx(l.Px()+fRand->Gaus(0,l.Px()*0.01));
    l.SetPy(l.Py()+fRand->Gaus(0,l.Py()*0.01));
    l.SetPz(l.Pz()+fRand->Gaus(0,l.Pz()*0.01));
    t.setP4(l);
  */
}

bool 
PndFastSim::cutAndSmear(PndFsmTrack *t)
{
  return cutAndSmear(t, t->detResponse());
}

//-----------------------------------------------------------------------

bool 
PndFastSim::cutAndSmear(PndFsmTrack *t, PndFsmResponse *r)
{
  if( !(r->detected()) ) return false;

  double charge = t->charge();

  double dE     = r->dE();
  double dp     = r->dp();
  double dtheta = r->dtheta();
  double dphi   = r->dphi();
  double dm     = r->dm();
  double m2     = r->m2();
  double MvddEdx =r->MvddEdx();
  double TpcdEdx =r->TpcdEdx();
  double SttdEdx =r->SttdEdx();
  TVector3 dV   = r->dV();
  

  //this removes candidates, which only have hit a PID-only device (like Cherenkov, TOF ...)
  
  if (fabs(charge)>1e-6 && fabs(dp)<1e-8) return false;
  if (fabs(charge)<1e-6 && fabs(dE)<1e-8) return false;
  
  if (dE != 0.0)     smearEnergy(t,dE);
  if (dp != 0.0)     smearMomentum(t,dp);
  if (dtheta != 0.0) smearTheta(t,dtheta);
  if (dphi != 0.0)   smearPhi(t,dphi);
  if (dm != 0.0)     smearM(t,dm);
  if( m2!=0.0)       smearM2(t,m2);           // mass^2 of track after tof
  if(MvddEdx!=0.0)   smearMvddEdx(t,MvddEdx); // dEdx of track after Mvd
  if(TpcdEdx!=0.0)   smearTpcdEdx(t,TpcdEdx); // dEdx of track after Tpc
  if(SttdEdx!=0.0)   smearSttdEdx(t,SttdEdx); // dEdx of track after Stt
  if (dV.X() != 0.0 || dV.Y() != 0.0 || dV.Z() != 0.0) smearVertex(t,dV);
  
  return true;
}

//-----------------------------------------------------------------------

void  
PndFastSim::smearEnergy(PndFsmTrack *t, double dE)
{
  TLorentzVector p4=t->p4();
  double m=t->p4().M();
  
  //  if (t->pdt()->lundId()==22) //gammas are always on mass shell
    if (fabs(t->charge()) < 0.1 ) //gammas are always on mass shell
  {
    //    double rnd = pRand->Gaus(0.,dE);
    double rnd = fRand->Gaus(0. , dE);

    double newE = rnd + p4.E();
    
    p4.SetVectM(p4.Vect()*(newE/p4.E()) , 0.0);
  }
  else
  {
    double newE = fRand->Gaus(0.0,dE) + p4.E();
    double newP = sqrt(newE*newE - m*m);

    p4.SetE(newE);
    p4.SetVectMag(p4.Vect(),newP);
  }
  
  t->setP4(p4);
}

void  
PndFastSim::smearMomentum(PndFsmTrack *t, double dp)
{
  TLorentzVector p4=t->p4();
  double newP = p4.Vect().Mag() + fRand->Gaus(0.0,dp);
  p4.SetVectM(p4.Vect()*(newP/p4.Vect().Mag()),t->p4().M());
  t->setP4(p4);
}

void  
PndFastSim::smearTheta(PndFsmTrack *t, double dtheta)
{
  TLorentzVector p4=t->p4();
  double newTheta = p4.Theta() + fRand->Gaus(0.,dtheta);
  p4.SetTheta(newTheta);
  
  t->setP4(p4);
}

void  
PndFastSim::smearPhi(PndFsmTrack *t, double dphi)
{
  TLorentzVector p4=t->p4();
  double newPhi = p4.Phi() + fRand->Gaus(0.,dphi); 
  p4.SetPhi(newPhi);
  
  t->setP4(p4);
}

void
PndFastSim::smearM(PndFsmTrack *t, double dm)
{
  TLorentzVector p4=t->p4();
//  double newM = p4.m() + RandGauss::shoot(0.,dm);
  double newM=p4.M() + dm;

  p4.SetVectM(p4.Vect(),newM);

  t->setP4(p4);
}

void
PndFastSim::smearM2(PndFsmTrack *t, double m2)
{

  t->setMass2(m2);
}

void
PndFastSim::smearMvddEdx(PndFsmTrack *t, double MvddEdx)
{

  t->setMvddEdX(MvddEdx);
}

void
PndFastSim::smearTpcdEdx(PndFsmTrack *t, double TpcdEdx)
{

  t->setTpcdEdX(TpcdEdx);
}

void
PndFastSim::smearSttdEdx(PndFsmTrack *t, double SttdEdx)
{

  t->setSttdEdX(SttdEdx);
}

void
PndFastSim::smearVertex(PndFsmTrack *t, TVector3 dV)
{
  TVector3 vtx = t->startVtx();

  double newX = vtx.X() + fRand->Gaus( 0. , dV.X() );
  double newY = vtx.Y() + fRand->Gaus( 0. , dV.Y() );
  double newZ = vtx.Z() + fRand->Gaus( 0. , dV.Z() );

  vtx.SetX(newX);
  vtx.SetY(newY);
  vtx.SetZ(newZ);

  t->setStartVtx(vtx);
}

//-----------------------------------------------------------------------


PndFsmResponse* 
PndFastSim::sumResponse(FsmResponseList respList)
{
  
  bool detected=false;

  double dE=0.0;
  double dp=0.0;
  double dtheta=0.0;
  double dphi=0.0;
  double dt=0.0;
  double dm=0.0;
  double m2=0;
  double MvddEdx=0;
  double TpcdEdx=0;
  double SttdEdx=0;
  double DrcDiscThtc=0;
  double DrcBarrelThtc=0;
  double RichThtc=0;
  double dVx=0.0;
  double dVy=0.0;
  double dVz=0.0;
  
  double LH_e=1.0;
  double LH_mu=1.0;
  double LH_pi=1.0;
  double LH_K=1.0;
  double LH_p=1.0;
  
  double val=0.0;
  
  PndFsmResponse *allResponse=new PndFsmResponse();
  
  if (fVb>3) cout <<" *** SINGLE RESPONSES ***"<<endl<<"--------------------------------------"<<endl;
  
  for (FsmResponseList::iterator riter=respList.begin(); riter!=respList.end();riter++)
  {
    PndFsmResponse *resp=*riter;
    
    if ( fVb>3 ) resp->print(cout);
    
    detected = detected | resp->detected();
    
    if (resp->detected())
    {
      if (fabs(val = resp->dE()) > 1e-8)     dE += 1./(val*val);
      if (fabs(val = resp->dp()) > 1e-8)      dp += 1./(val*val);
      if (fabs(val = resp->dtheta())> 1e-8) dtheta += 1./(val*val);
      if (fabs(val = resp->dphi()) > 1e-8)   dphi += 1./(val*val);
      if (fabs(val = resp->dt()) > 1e-8)     dt += val*val;
      if (fabs(val = resp->dm()) > 1e-8)     dm +=val;
      if (fabs (val = resp->m2()) > 1e-11)    m2+=val;
      if (fabs (val = resp->MvddEdx()) > 1e-11)    MvddEdx+=val;
      if (fabs (val = resp->TpcdEdx()) > 1e-11)    TpcdEdx+=val;
      if (fabs (val = resp->SttdEdx()) > 1e-11)    SttdEdx+=val;
      if (fabs (val = resp->DrcDiscThtc()) > 1e-11)   DrcDiscThtc+=val;
      if (fabs (val = resp->DrcBarrelThtc()) > 1e-11)    DrcBarrelThtc+=val;
      if (fabs (val = resp->RichThtc()) > 1e-11)    RichThtc+=val;

      if (fabs (val = resp->dV().X()) > 1e-11) dVx += 1./(val*val);
      if (fabs (val = resp->dV().Y()) > 1e-11) dVy += 1./(val*val);
      if (fabs (val = resp->dV().Z()) > 1e-11) dVz += 1./(val*val);

      double rawLHe  = resp->LHElectron();
      double rawLHmu = resp->LHMuon();
      double rawLHpi = resp->LHPion();
      double rawLHK  = resp->LHKaon();
      double rawLHp  = resp->LHProton();

      double sumRaw = rawLHe+rawLHmu+rawLHpi+rawLHK+rawLHp;

      if (sumRaw>0) {
	rawLHe  /= sumRaw;
	rawLHmu /= sumRaw;
	rawLHpi /= sumRaw;
	rawLHK  /= sumRaw;
	rawLHp  /= sumRaw;
      }

      //here a weighted Likelihood evaluation has to be done
      LH_e  *= rawLHe; 
      LH_mu *= rawLHmu; 
      LH_pi *= rawLHpi; 
      LH_K  *= rawLHK; 
      LH_p  *= rawLHp; 

      /*
      if (val = rawLHe)  LH_e  == 0.0 ? LH_e  = val : LH_e  *= val; 
      if (val = rawLHmu) LH_mu == 0.0 ? LH_mu = val : LH_mu *= val; 
      if (val = rawLHpi) LH_pi == 0.0 ? LH_pi = val : LH_pi *= val; 
      if (val = rawLHK)  LH_K  == 0.0 ? LH_K  = val : LH_K  *= val; 
      if (val = rawLHp)  LH_p  == 0.0 ? LH_p  = val : LH_p  *= val; 
      */
    } 
  }
  
  double sumLH = LH_e + LH_mu + LH_pi + LH_K + LH_p;

  LH_e  /= sumLH;
  LH_mu /= sumLH;
  LH_pi /= sumLH;
  LH_K  /= sumLH;
  LH_p  /= sumLH;
  
  allResponse->setDetected(detected);
  allResponse->setdE( dE>0. ? sqrt(1./dE) : 0.0 );
  allResponse->setdp( dp>0. ? sqrt(1./dp) : 0.0 );
  allResponse->setdtheta( dtheta>0. ? sqrt(1./dtheta) : 0.0 );
  allResponse->setdphi( dphi>0. ? sqrt(1./dphi) : 0.0 );
  allResponse->setdt( sqrt(dt) );
  allResponse->setdm(dm);
  allResponse->setm2(m2);
  allResponse->setMvddEdx(MvddEdx);
  allResponse->setTpcdEdx(TpcdEdx);
  allResponse->setSttdEdx(SttdEdx);
  
  allResponse->setDrcDiscThtc(DrcDiscThtc);
  allResponse->setDrcBarrelThtc(DrcBarrelThtc);
  allResponse->setRichThtc(RichThtc);
  
  if (dVx > 0.) dVx=1./sqrt(dVx); else dVx = 0.0;
  if (dVy > 0.) dVy=1./sqrt(dVy); else dVy = 0.0;
  if (dVz > 0.) dVz=1./sqrt(dVz); else dVz = 0.0;
  
  allResponse->setdV( dVx , dVy , dVz );
  
  allResponse->setLHElectron(LH_e);
  allResponse->setLHMuon(LH_mu);
  allResponse->setLHPion(LH_pi);
  allResponse->setLHKaon(LH_K);
  allResponse->setLHProton(LH_p);	  
  
  
  if (fVb>2) cout <<"--------------------------------------"<<endl<<"*** OVERALL RESPONSE ***"<<endl<<"--------------------------------------" <<endl;
  if (fVb>2) allResponse->print(cout);
  if (fVb>2) cout <<"--------------------------------------"<<endl;
  
  return allResponse; 
}

//-----------------------------------------------------------------------

ClassImp(PndFastSim)