#include "modules/RunPUPPI.h"

#include "PUPPI/RecoObj2.hh"
#include "PUPPI/AlgoObj.hh"
#include "PUPPI/PuppiContainer.hh"

#include "fastjet/PseudoJet.hh"

#include "classes/DelphesClasses.h"
#include "classes/DelphesFactory.h"
#include "classes/DelphesFormula.h"

#include <algorithm> 
#include <stdexcept>
#include <iostream>
#include <sstream>
#include <vector>

using namespace std;
using namespace fastjet;

//------------------------------------------------------------------------------
RunPUPPI::RunPUPPI() :
  fItTrackInputArray(0), 
  fItNeutralInputArray(0)
{}

//------------------------------------------------------------------------------
RunPUPPI::~RunPUPPI(){}

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

void RunPUPPI::Init(){
  // input collection
  fTrackInputArray     = ImportArray(GetString("TrackInputArray", "Calorimeter/towers"));
  fItTrackInputArray   = fTrackInputArray->MakeIterator();
  fNeutralInputArray   = ImportArray(GetString("NeutralInputArray", "Calorimeter/towers"));
  fItNeutralInputArray = fNeutralInputArray->MakeIterator();
  fPVInputArray        = ImportArray(GetString("PVInputArray", "PV"));
  fPVItInputArray      = fPVInputArray->MakeIterator();
  // puppi parameters                                 
  fApplyNoLep     = GetBool("UseNoLep", true);    
  fMinPuppiWeight = GetDouble("MinPuppiWeight", 0.01);
  fUseExp         = GetBool("UseExp", false);
  // read eta min ranges
  ExRootConfParam param = GetParam("EtaMinBin");
  fEtaMinBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fEtaMinBin.push_back(param[iMap].GetDouble());
  // read eta max ranges
  param = GetParam("EtaMaxBin");
  fEtaMaxBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fEtaMaxBin.push_back(param[iMap].GetDouble());
  // read pt min value
  param = GetParam("PtMinBin");
  fPtMinBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fPtMinBin.push_back(param[iMap].GetDouble());
  // read cone size                                                                                                                                                           
  param = GetParam("ConeSizeBin");
  fConeSizeBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fConeSizeBin.push_back(param[iMap].GetDouble());
  // read RMS min pt                                                                                                                                             
  param = GetParam("RMSPtMinBin");
  fRMSPtMinBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fRMSPtMinBin.push_back(param[iMap].GetDouble());
  // read RMS scale factor
  param = GetParam("RMSScaleFactorBin");
  fRMSScaleFactorBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fRMSScaleFactorBin.push_back(param[iMap].GetDouble());
  // read neutral pt min cut
  param = GetParam("NeutralMinEBin");
  fNeutralMinEBin.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fNeutralMinEBin.push_back(param[iMap].GetDouble());
  // read neutral pt min slope
  param = GetParam("NeutralPtSlope");
  fNeutralPtSlope.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fNeutralPtSlope.push_back(param[iMap].GetDouble());
  // read apply chs                                                                                                                                                           
  //param = GetParam("ApplyCHS");
  //fApplyCHS.clear();
  //for(int iMap = 0; iMap < param.GetSize(); ++iMap) fApplyCHS.push_back(param[iMap].GetBool());
  // read use charged
  param = GetParam("UseCharged");
  fUseCharged.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fUseCharged.push_back(param[iMap].GetBool());
  // read apply chs correction
  param = GetParam("ApplyLowPUCorr");
  fApplyLowPUCorr.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fApplyLowPUCorr.push_back(param[iMap].GetBool());
  // read metric id                                                                                                                                                          
  param = GetParam("MetricId");
  fMetricId.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fMetricId.push_back(param[iMap].GetInt());
  // scheme for combining
  param = GetParam("CombId");
  fCombId.clear();
  for(int iMap = 0; iMap < param.GetSize(); ++iMap) fCombId.push_back(param[iMap].GetInt());
  // create output array
  fOutputArray        = ExportArray(GetString("OutputArray", "puppiParticles"));
  fOutputTrackArray   = ExportArray(GetString("OutputArrayTracks", "puppiTracks"));
  fOutputNeutralArray = ExportArray(GetString("OutputArrayNeutrals", "puppiNeutrals"));
  // Create algorithm list for puppi
  std::vector<AlgoObj> puppiAlgo;
  if(puppiAlgo.empty()){
    if(!(fEtaMinBin.size() == fEtaMaxBin.size() and fEtaMinBin.size() == fPtMinBin.size() and fEtaMinBin.size() == fConeSizeBin.size() and fEtaMinBin.size() == fRMSPtMinBin.size()
	 and fEtaMinBin.size() == fRMSScaleFactorBin.size() and fEtaMinBin.size() == fNeutralMinEBin.size() and  fEtaMinBin.size() == fNeutralPtSlope.size() 
	 and fEtaMinBin.size() == fUseCharged.size()
	 and fEtaMinBin.size() == fApplyLowPUCorr.size() and fEtaMinBin.size() == fMetricId.size())) {
      std::cerr<<" Error in PUPPI configuration, algo info should have the same size --> exit from the code"<<std::endl;
      std::exit(EXIT_FAILURE);
    }
  }
  for( size_t iAlgo =  0 ; iAlgo < fEtaMinBin.size() ; iAlgo++){
    AlgoObj algoTmp ;
    algoTmp.etaMin            = fEtaMinBin.at(iAlgo);
    algoTmp.etaMax            = fEtaMaxBin.at(iAlgo);
    algoTmp.ptMin             = fPtMinBin.at(iAlgo);
    algoTmp.minNeutralPt      = fNeutralMinEBin.at(iAlgo);
    algoTmp.minNeutralPtSlope = fNeutralPtSlope.at(iAlgo);
    //Eta Extrapolation stuff is missing
    //Loop through file requiring algos for same bins to be adjacent
    while(iAlgo < fEtaMinBin.size() and algoTmp.etaMin == fEtaMinBin.at(iAlgo) and algoTmp.etaMax == fEtaMaxBin.at(iAlgo)) { 
      AlgoSubObj algoSubTmp;
      algoSubTmp.metricId          = fMetricId.at(iAlgo);
      algoSubTmp.useCharged        = fUseCharged.at(iAlgo);
      algoSubTmp.applyLowPUCorr    = fApplyLowPUCorr.at(iAlgo);
      algoSubTmp.combId            = fCombId.at(iAlgo);
      algoSubTmp.coneSize          = fConeSizeBin.at(iAlgo);
      algoSubTmp.rmsPtMin          = fRMSPtMinBin.at(iAlgo);
      algoSubTmp.rmsScaleFactor    = fRMSScaleFactorBin.at(iAlgo);
      algoTmp.subAlgos.push_back(algoSubTmp);
      iAlgo++;
    }
    iAlgo--;
    //if(std::find(puppiAlgo.begin(),puppiAlgo.end(),algoTmp) != puppiAlgo.end()) continue;    
    puppiAlgo.push_back(algoTmp);     
  }
  fPuppi  = new PuppiContainer(true,fUseExp,fMinPuppiWeight,puppiAlgo);
}

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

void RunPUPPI::Finish(){
  if(fItTrackInputArray)   delete fItTrackInputArray;
  if(fItNeutralInputArray) delete fItNeutralInputArray;
}

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

void RunPUPPI::Process(){

  Candidate *candidate, *particle;
  TLorentzVector momentum;

  //DelphesFactory *factory = GetFactory();

  // loop over input objects
  fItTrackInputArray   ->Reset();
  fItNeutralInputArray ->Reset();
  fPVItInputArray      ->Reset();

  std::vector<Candidate *> InputParticles;
  InputParticles.clear();

  // take the leading vertex 
  float PVZ = 0.;
  Candidate *pv = static_cast<Candidate*>(fPVItInputArray->Next());
  if (pv) PVZ = pv->Position.Z();
  // Fill input particles for puppi
  std::vector<RecoObj> puppiInputVector;
  puppiInputVector.clear();
  int lNBad  = 0; 
  // Loop on charge track candidate
  while((candidate = static_cast<Candidate*>(fItTrackInputArray->Next()))){   
      momentum = candidate->Momentum;
      RecoObj curRecoObj;
      curRecoObj.pt  = momentum.Pt();
      curRecoObj.eta = momentum.Eta();
      curRecoObj.phi = momentum.Phi();
      curRecoObj.m   = momentum.M();  
      particle = static_cast<Candidate*>(candidate->GetCandidates()->At(0));//if(fApplyNoLep && TMath::Abs(candidate->PID) == 11) continue; //Dumb cut to minimize the nolepton on electron
      //if(fApplyNoLep && TMath::Abs(candidate->PID) == 13) continue;
      if (candidate->IsRecoPU and candidate->Charge !=0) { // if it comes fromPU vertexes after the resolution smearing and the dZ matching within resolution
	lNBad++;
	curRecoObj.id    = 2;
	curRecoObj.vtxId = 0.7*(fPVInputArray->GetEntries()); //Hack apply reco vtx efficiency of 70% for calibration
	if(TMath::Abs(candidate->PID) == 11)      curRecoObj.pfType = 2;
	else if(TMath::Abs(candidate->PID) == 13) curRecoObj.pfType = 3;
	else if(TMath::Abs(candidate->PID) == 22) curRecoObj.pfType = 4;
        else curRecoObj.pfType = 1;
        curRecoObj.dZ = particle->Position.Z()-PVZ;
      } 
      else if(!candidate->IsRecoPU && candidate->Charge !=0) {
	curRecoObj.id    = 1;  // charge from LV
        curRecoObj.vtxId = 1; // from PV
	if(TMath::Abs(candidate->PID) == 11)      curRecoObj.pfType = 2;
	else if(TMath::Abs(candidate->PID) == 13) curRecoObj.pfType = 3;
	else if(TMath::Abs(candidate->PID) == 22) curRecoObj.pfType = 4;
        else curRecoObj.pfType = 1;
        curRecoObj.dZ = particle->Position.Z()-PVZ;
      }
      else {
	std::cerr<<" RunPUPPI: problem with a charged track --> it has charge 0 "<<std::endl;
        continue;
      }

      puppiInputVector.push_back(curRecoObj);
      InputParticles.push_back(candidate);
  }

  // Loop on neutral calo cells 
  while((candidate = static_cast<Candidate*>(fItNeutralInputArray->Next()))){
      momentum = candidate->Momentum;
      RecoObj curRecoObj;
      curRecoObj.pt  = momentum.Pt();
      curRecoObj.eta = momentum.Eta();
      curRecoObj.phi = momentum.Phi();
      curRecoObj.m   = momentum.M();
      curRecoObj.charge = 0;
      particle = static_cast<Candidate*>(candidate->GetCandidates()->At(0));
      if(candidate->Charge == 0){
	curRecoObj.id    = 0; // neutrals have id==0
	curRecoObj.vtxId = 0; // neutrals have vtxId==0 
	if(TMath::Abs(candidate->PID) == 11)      curRecoObj.pfType = 2;
	else if(TMath::Abs(candidate->PID) == 13) curRecoObj.pfType = 3;
	else if(TMath::Abs(candidate->PID) == 22) curRecoObj.pfType = 4;
        else curRecoObj.pfType = 5;
        curRecoObj.dZ = particle->Position.Z()-PVZ;
      }
      else{
	std::cerr<<" RunPUPPI: problem with a neutrals cells --> it has charge !=0 "<<std::endl;
        continue;
      }
      puppiInputVector.push_back(curRecoObj);
      InputParticles.push_back(candidate);
  }
  // Create PUPPI container
  fPuppi->initialize(puppiInputVector);
  fPuppi->puppiWeights();
  std::vector<PseudoJet> puppiParticles = fPuppi->puppiParticles();

  // Loop on final particles
  for (std::vector<PseudoJet>::iterator it = puppiParticles.begin() ; it != puppiParticles.end() ; it++) {
    if(it->user_index() <= int(InputParticles.size())){      
      candidate = static_cast<Candidate *>(InputParticles.at(it->user_index())->Clone());
      candidate->Momentum.SetPxPyPzE(it->px(),it->py(),it->pz(),it->e());
      fOutputArray->Add(candidate);
      if( puppiInputVector.at(it->user_index()).id == 1 or puppiInputVector.at(it->user_index()).id == 2) fOutputTrackArray->Add(candidate);
      else if (puppiInputVector.at(it->user_index()).id == 0) fOutputNeutralArray->Add(candidate);
    }
    else{ 
      std::cerr<<" particle not found in the input Array --> skip "<<std::endl;
      continue;
    }     
  }

}