#include <iostream>
#include <vector>
#include <string>
#include <cstring>
#include <iomanip>
#include <cmath>

#include <HEPUnits.hxx>
#include <HEPConstants.hxx>
#include <IRealDatum.hxx>

#include "IPIDManager.hxx"

#include "IOARuntimeParameters.hxx"


 /// The static member pointer to the singleton.
IPIDManager* IPIDManager::fTPIDManager = NULL;


//*****************************************************************************
IPIDManager& COMET::pman(){
//*****************************************************************************

  return IPIDManager::Get();

}

//*****************************************************************************
IPIDManager::IPIDManager() {
//*****************************************************************************

  SetVerbosity(0);

  fLastPID = NULL;
  fComputedProbs = false;

  fPidHypotheses.push_back(COMET::IReconPID::kElectron);
  fPidHypotheses.push_back(COMET::IReconPID::kMuon);
  fPidHypotheses.push_back(COMET::IReconPID::kProton);
}

//*****************************************************************************
IPIDManager& IPIDManager::Get(void) {
//*****************************************************************************
  
  if (!fTPIDManager) {
    std::cout << "Create a new PIDManager object."
              << std::endl;
    fTPIDManager = new IPIDManager;
  }
  

  return *fTPIDManager;
}


//*****************************************************************************
const LikeMap& IPIDManager::GetLikelihoodMap(const COMET::IHandle<COMET::IReconPID> pid){
//*****************************************************************************
  
  if (GetPointer(pid) == fLastPID)
    return fLikes;

  fLikes.clear();

  // save main hypothesis
  fLikes[pid->GetParticleId()] = pid->GetPIDWeight();

  // save alternate hypothesis
  const COMET::IReconObjectContainer&  alter = pid->GetAlternates();
  COMET::IReconObjectContainer::const_iterator it;  
  for (it = alter.begin(); it!=alter.end(); it++){
    COMET::IHandle<COMET::IReconPID> apid = *it;
    fLikes[apid->GetParticleId()] = apid->GetPIDWeight();
  }

  fComputedProbs = false;

  return fLikes;

}

//*****************************************************************************
const ProbMap& IPIDManager::GetProbabilityMap(const COMET::IHandle<COMET::IReconPID> pid){
//*****************************************************************************

  if (GetPointer(pid) == fLastPID && fComputedProbs)
    return fProbs;
  
  fLikes = GetLikelihoodMap(pid);

  fProbs.clear();
  fSortedProbs.clear();

  double sum=0;
  LikeMap::const_iterator it;
  for (it = fLikes.begin(); it!=fLikes.end(); it++){
    sum += it->second; 
  }

  for (it = fLikes.begin(); it!=fLikes.end(); it++){
    if (sum != 0){
      fProbs[it->first] = fLikes[it->first]/sum;    
      fSortedProbs[fProbs[it->first]] = it->first; 
    }
    else{
      fProbs[it->first] = 0;
    }
  }

  fComputedProbs = true;

  return fProbs;

}

//*****************************************************************************
const ProbMap& IPIDManager::GetProbabilityMap(const COMET::IHandle<COMET::IReconPID> pid, 
                                              const PriorMap& priors){
//*****************************************************************************

  fProbs = GetProbabilityMap(pid);

  fPriorProbs.clear();
  fSortedPriorProbs.clear();

  double sum=0;
  ProbMap::const_iterator it;
  for (it = fProbs.begin(); it!=fProbs.end(); it++){
    if (priors.find(it->first) != priors.end())
      sum += fProbs[it->first] * priors.find(it->first)->second; 
  }

  for (it = fProbs.begin(); it!=fProbs.end(); it++){
    if (sum != 0 &&  priors.find(it->first) != priors.end()){      
        fPriorProbs[it->first] = (fProbs[it->first]* priors.find(it->first)->second)/sum;    
        fSortedPriorProbs[fPriorProbs[it->first]] = it->first; 
    }
    else{
      fPriorProbs[it->first] = 0;
    }
  }


  return fPriorProbs;

}

//*****************************************************************************
const SortedProbMap& IPIDManager::GetSortedProbabilityMap(const COMET::IHandle<COMET::IReconPID> pid){
//*****************************************************************************

  fProbs = GetProbabilityMap(pid);

  return fSortedProbs;
  
}


//*****************************************************************************
const SortedProbMap& IPIDManager::GetSortedProbabilityMap(const COMET::IHandle<COMET::IReconPID> pid, const PriorMap& priors){
//*****************************************************************************

  fPriorProbs = GetProbabilityMap(pid,priors);

  return fSortedPriorProbs;
  
}



//*****************************************************************************
double IPIDManager::GetProbability(const COMET::IHandle<COMET::IReconPID> pid, 
                                   COMET::IReconPID::ParticleId id){
//*****************************************************************************

  if (GetProbabilityMap(pid).find(id) != GetProbabilityMap(pid).end())
    return GetProbabilityMap(pid).find(id)->second;
  else 
    return 0;
}


//*****************************************************************************
void IPIDManager::DumpProbabilities(const COMET::IHandle<COMET::IReconPID> pid){
  //*****************************************************************************
  
  fSortedProbs = GetSortedProbabilityMap(pid);

  std::map<double, COMET::IReconPID::ParticleId>::reverse_iterator it;
  for (it = fSortedProbs.rbegin(); it!=fSortedProbs.rend(); it++){
    std::cout << " - hypo = " << it->second << " (" << COMET::IReconPID::ConvertParticleId(it->second) << ")" 
              << "\t ---> prob = " << it->first 
              << std::endl;
  }

}


//*****************************************************************************
void IPIDManager::DumpProbabilities(const COMET::IHandle<COMET::IReconPID> pid, const PriorMap& priors){
  //*****************************************************************************
  
  fSortedPriorProbs = GetSortedProbabilityMap(pid,priors);

  SortedProbMap::reverse_iterator it;
  for (it = fSortedPriorProbs.rbegin(); it!=fSortedPriorProbs.rend(); it++){
    std::cout << " - hypo = " << it->second << " (" << COMET::IReconPID::ConvertParticleId(it->second) << ")" 
              << "\t ---> prob = " << fProbs[it->second] 
              << ", prior = " << priors.find(it->second)->second 
              << ", final prob = " << it->first
              << std::endl;
  }

}


//*****************************************************************************
void IPIDManager::CopyPIDInfo(COMET::IHandle<COMET::IReconPID> PID1, COMET::IHandle<COMET::IReconPID> PID2){
//*****************************************************************************

  if (GetVerbosity()>2){  
    std::cout << "IPIDManager::CopyPIDInfo()."  << GetPointer(PID1) << " --> " << GetPointer(PID2) << std::endl;
  }

  /*
  PID2->SetParticleId(PID1->GetParticleId());
  PID2->SetPIDWeight(PID1->GetPIDWeight());

  const COMET::IReconObjectContainer&  pids1 = PID1->GetAlternates();

  COMET::IReconObjectContainer::const_iterator it;  
  for (it = pids1.begin(); it!=pids1.end(); it++){
    COMET::IHandle<COMET::IReconPID> pid = *it;
    PID2->AddAlternate(pid->GetParticleId(), pid->GetPIDWeight() );
  }
  */

  // copy the PID weights
  CopyPIDWeightsFromIDatums(PID1, PID2);

  std::map<double, COMET::IReconPID::ParticleId> w2map;  

  // get the weights 
  double w2[5];
  GetPIDWeightsFromIDatums(PID2, w2);

  // same them in a map used to get the maximum weight
  for (unsigned int i=0;i<fPidHypotheses.size();i++){
    w2map[w2[i]] = fPidHypotheses[i]; 
  }

  // take the hypothesis with the maximum weight
  PID2->SetParticleId(w2map.rbegin()->second);
  PID2->SetPIDWeight(w2map.rbegin()->first);


}

//*****************************************************************************
bool IPIDManager::MergePIDInfo(COMET::IHandle<COMET::IReconBase> object1, COMET::IHandle<COMET::IReconBase> object2,
                               COMET::IHandle<COMET::IReconPID> PID3){
//*****************************************************************************

  COMET::IHandle<COMET::IReconPID> PID1 = object1;
  COMET::IHandle<COMET::IReconPID> PID2 = object2;

  COMET::IReconObjectContainer::const_iterator it;  

  if (GetVerbosity()>2){
    
    std::cout << "IPIDManager::MergePIDInfo()." << std::endl;
    std::cout << " - FirstObject: " << GetPointer(object1) << std::endl;
    if (PID1){
      std::cout << "   - PID = " << PID1->ConvertParticleId() << ", w = " << PID1->GetPIDWeight() << std::endl;    
      const COMET::IReconObjectContainer&  pids1 = PID1->GetAlternates();
      for (it = pids1.begin(); it!=pids1.end(); it++){
        COMET::IHandle<COMET::IReconPID> pid = *it;
        std::cout << "   - PID = " << pid->ConvertParticleId() << ", w = " << pid->GetPIDWeight() << std::endl;
      }
    }

    std::cout << " - SecondObject: " << GetPointer(object2) << std::endl;
    if (PID2){
      std::cout << "   - PID = " << PID2->ConvertParticleId() << ", w = " << PID2->GetPIDWeight() << std::endl;    
      const COMET::IReconObjectContainer&  pids2 = PID2->GetAlternates();
      for (it = pids2.begin(); it!=pids2.end(); it++){
        COMET::IHandle<COMET::IReconPID> pid = *it;
        std::cout << "   - PID = " << pid->ConvertParticleId() << ", w = " << pid->GetPIDWeight() << std::endl;
      }
    }

  }

  // ------ Check that objects are TReconPIDs ------------ 

  // 1. none of the objects is a IReconPID
  if (!PID1 && !PID2) return false;
  // 2. First object is a IReconPID
  else if (PID1 && !PID2){
    if (PID1->GetParticleId() != COMET::IReconPID::kNotSet &&  CheckPIDDetector(PID1) ){
      CopyPIDInfo(PID1, PID3);
      return true;
    }
    return false;
  }
  // 3. Second object is a IReconPID
  else if (!PID1 && PID2){
    if (PID2->GetParticleId() != COMET::IReconPID::kNotSet &&  CheckPIDDetector(PID2) ){
      CopyPIDInfo(PID2, PID3);
      return true;
    }
    return false;
  }
    
  // 4. Both objects are  TReconPIDs
  

  // ------ Check that objects have a valid PID hypothesis ----------- -

  // 4.1. Both objects have no PID hypothesis or are not used in combine PID
  if ((PID1->GetParticleId() == COMET::IReconPID::kNotSet || !CheckPIDDetector(PID1) ) && 
      (PID2->GetParticleId() == COMET::IReconPID::kNotSet || !CheckPIDDetector(PID2) )) 
    return false;
  // 4.2. First object have PID hypothesis and is used in combinen PID
  else if ( PID1->GetParticleId() != COMET::IReconPID::kNotSet  &&  CheckPIDDetector(PID1) && 
            (PID2->GetParticleId() == COMET::IReconPID::kNotSet || !CheckPIDDetector(PID2) ) ){
    CopyPIDInfo(PID1, PID3);
    return true;
  }
  // 4.3. Second object have PID hypothesis and is used in combinen PID
  else if ((PID1->GetParticleId() == COMET::IReconPID::kNotSet || !CheckPIDDetector(PID1) ) && 
           PID2->GetParticleId() != COMET::IReconPID::kNotSet  &&  CheckPIDDetector(PID2)){
    CopyPIDInfo(PID2, PID3);
    return true;
  }
  
  // 4.4. Both objects have PID hypothesis and are used in combined PID 

  // ------ Merge PID info ------------

  const COMET::IReconObjectContainer&  pids1 = PID1->GetAlternates();
  const COMET::IReconObjectContainer&  pids2 = PID2->GetAlternates();

  std::map<COMET::IReconPID::ParticleId, std::pair<double, double> > weights;  

  // 1.  save main PID hypothesis in the map
  if (PID1->GetParticleId() == PID2->GetParticleId()){
    std::pair<double, double> w(PID1->GetPIDWeight(),PID2->GetPIDWeight());
    weights[PID1->GetParticleId()] = w;
  }
  else{
    std::pair<double, double> w1(PID1->GetPIDWeight(),0);
    weights[PID1->GetParticleId()] = w1;

    std::pair<double, double> w2(0,PID2->GetPIDWeight());
    weights[PID2->GetParticleId()] = w2;

  }

  if (GetVerbosity()>2){
    std::cout << "IPIDManager::MergePIDInfo(). weights afters main hypothesis:" << std::endl;      
    
    std::map<COMET::IReconPID::ParticleId, std::pair<double, double> >::const_iterator it2;
    for (it2 =weights.begin(); it2 != weights.end(); it2++){
      if (GetVerbosity()>2){    
        std::cout << " - " << COMET::IReconPID::ConvertParticleId(it2->first) << ": " 
                  << (it2->second.first) << ", " << (it2->second.second) << std::endl;
      }
    }
  }

  // 2. save alternates in the map
  for (it = pids1.begin(); it!=pids1.end(); it++){
    COMET::IHandle<COMET::IReconPID> pid = *it;
    std::pair<double, double> w;
    if (weights.find(pid->GetParticleId()) == weights.end()) 
      w = std::pair<double, double>(pid->GetPIDWeight(),0);
    else{
      if (weights[pid->GetParticleId()].first == 0)
        w = std::pair<double, double>(pid->GetPIDWeight(), weights[pid->GetParticleId()].second );
      else
        w = std::pair<double, double>(weights[pid->GetParticleId()].first, pid->GetPIDWeight() );
    }

    weights[pid->GetParticleId()] = w;
  }

  if (GetVerbosity()>2){
    std::cout << "IPIDManager::MergePIDInfo(). weights afters first alternates:" << std::endl;      
    
    std::map<COMET::IReconPID::ParticleId, std::pair<double, double> >::const_iterator it2;
    for (it2 =weights.begin(); it2 != weights.end(); it2++){
      if (GetVerbosity()>2){    
        std::cout << " - " << COMET::IReconPID::ConvertParticleId(it2->first) << ": " 
                  << (it2->second.first) << ", " << (it2->second.second) << std::endl;
      }
    }
  }


  for (it = pids2.begin(); it!=pids2.end(); it++){
    COMET::IHandle<COMET::IReconPID> pid = *it;
    std::pair<double, double> w;
    if (weights.find(pid->GetParticleId()) == weights.end()) 
      w = std::pair<double, double>(0,pid->GetPIDWeight());
    else{
      if (weights[pid->GetParticleId()].first == 0)
        w = std::pair<double, double>(pid->GetPIDWeight(), weights[pid->GetParticleId()].second );
      else
        w = std::pair<double, double>(weights[pid->GetParticleId()].first, pid->GetPIDWeight() );
    }


    
    weights[pid->GetParticleId()] = w;
  }

  if (GetVerbosity()>2)
    std::cout << "IPIDManager::MergePIDInfo(). Compute final weights:" << std::endl;      
  

  // 3. compute final weights
  std::map<double, COMET::IReconPID::ParticleId > final_weights;  
  std::map<COMET::IReconPID::ParticleId, std::pair<double, double> >::const_iterator it2;
  double total_w = 0;
  for (it2 =weights.begin(); it2 != weights.end(); it2++){
    total_w += (it2->second.first)*(it2->second.second); 
  }
  for (it2 =weights.begin(); it2 != weights.end(); it2++){

    double w = (it2->second.first)*(it2->second.second)/total_w;

    if (GetVerbosity()>2){    
      std::cout << " - " << COMET::IReconPID::ConvertParticleId(it2->first) << ": " 
                << (it2->second.first) << " * " << (it2->second.second) << " = " << w << std::endl;
    }

    if (w > 0)
      final_weights[w] = it2->first;
  }


  // 4. save information in the merged PID
  std::map<double, COMET::IReconPID::ParticleId >::reverse_iterator it3;
  for (it3 = final_weights.rbegin(); it3 != final_weights.rend(); it3++){
    if (it3 == final_weights.rbegin()){ 
      PID3->SetParticleId(it3->second);
      PID3->SetPIDWeight(it3->first);
    }
    else
      PID3->AddAlternate( it3->second, it3->first);
  }

  if (GetVerbosity()>0){

    // 5. Dump info
    
    std::cout << " Merged PID results -------------------------- " << std::endl;
    
    std::cout << " - PID = " << PID3->ConvertParticleId() << ", w = " << PID3->GetPIDWeight() << std::endl;
    
    const COMET::IReconObjectContainer&  pids3 = PID3->GetAlternates();
    for (it = pids3.begin(); it!=pids3.end(); it++){
      COMET::IHandle<COMET::IReconPID> pid = *it;
      std::cout << " - PID = " << pid->ConvertParticleId() << ", w = " << pid->GetPIDWeight() << std::endl;
    }
  }

  return true;
}

//*****************************************************************************
void IPIDManager::AddPIDDetector(COMET::IReconBase::StateBits det){
//*****************************************************************************

  fPidDetectors.push_back(det);
}

//*****************************************************************************
bool IPIDManager::CheckPIDDetector(COMET::IHandle<COMET::IReconBase> object){
//*****************************************************************************

  std::vector< COMET::IReconBase::StateBits >::const_iterator it;
  for (it=fPidDetectors.begin(); it!=fPidDetectors.end(); it++){
    if (object->UsesDetector(*it)) return true;
  }
  
  return false;
}

//*****************************************************************************
void IPIDManager::AddPIDHypothesis(COMET::IReconPID::ParticleId hyp){
//*****************************************************************************

  fPidHypotheses.push_back(hyp);
}

//*****************************************************************************
std::string IPIDManager::GetPIDHypothesisForIDatum(COMET::IReconPID::ParticleId hyp){
//*****************************************************************************

  if      (hyp == COMET::IReconPID::kElectron) return "Electron";
  else if (hyp == COMET::IReconPID::kMuon)     return "Muon";
  else if (hyp == COMET::IReconPID::kPion)     return "Pion";
  else if (hyp == COMET::IReconPID::kKaon)     return "Kaon";
  else if (hyp == COMET::IReconPID::kProton)   return "Proton";
  else return "";
}

//*****************************************************************************
void IPIDManager::GetPIDWeightsFromIDatums(COMET::IHandle<COMET::IReconPID> pid, double w[]){
//*****************************************************************************

  if (GetVerbosity()>2){
    std::cout << "IPIDManager::GetPIDWeightsFromIDatums" << std::endl;
  }

  COMET::IHandle<COMET::IRealDatum> datum=pid->Get<COMET::IRealDatum>("WeightElectron");
  if (!datum) return;

  for (unsigned int i=0;i<fPidHypotheses.size();i++){
    std::string name = "Weight"+ GetPIDHypothesisForIDatum(fPidHypotheses[i]);
    w[i] = pid->Get<COMET::IRealDatum>(name.c_str())->GetValue();
  }

  if (GetVerbosity()>2){
    DumpPIDWeightsFromIDatums(pid);
  }
}

//*****************************************************************************
void IPIDManager::SavePIDWeightsInIDatums( double w[], COMET::IHandle<COMET::IReconPID> pid){
//*****************************************************************************

  if (GetVerbosity()>2){
    std::cout << "IPIDManager::SavePIDWeightsInIDatums" << std::endl;
  }

  for (unsigned int i=0;i<fPidHypotheses.size();i++){
    std::string name = "Weight"+ GetPIDHypothesisForIDatum(fPidHypotheses[i]);
    COMET::IHandle<COMET::IRealDatum> datum0 = pid->Get<COMET::IRealDatum>(name.c_str());
    if(datum0)
      datum0->SetValue(w[i]);
    else{
      COMET::IRealDatum* datum = new COMET::IRealDatum(name.c_str(), w[i]);
      pid->AddDatum(datum);
    }
  }

  if (GetVerbosity()>2){
    DumpPIDWeightsFromIDatums(pid);
  }
}

//*****************************************************************************
void IPIDManager::CopyPIDWeightsFromIDatums( COMET::IHandle<COMET::IReconBase> object1, COMET::IHandle<COMET::IReconBase> object2){
//*****************************************************************************

  COMET::IHandle<COMET::IReconPID> PID1 = object1;
  COMET::IHandle<COMET::IReconPID> PID2 = object2;
  
  if (!PID1 || !PID2) return;

  COMET::IHandle<COMET::IRealDatum> datum=PID1->Get<COMET::IRealDatum>("WeightElectron");
  if(datum){
    // Get the PID weights of PID1
    double w1[5];
    GetPIDWeightsFromIDatums(PID1,w1);

    // save the PID weight of PID1 in PID2
    SavePIDWeightsInIDatums(w1, PID2);
  }
}

//*****************************************************************************
void IPIDManager::DumpPIDWeightsFromIDatums(COMET::IHandle<COMET::IReconPID> PID){
//*****************************************************************************

  COMET::IHandle<COMET::IRealDatum> datum=PID->Get<COMET::IRealDatum>("WeightElectron");
  if (!datum) return;

  double w[5];
  for (unsigned int i=0;i<fPidHypotheses.size();i++){
    std::string name = "Weight"+ GetPIDHypothesisForIDatum(fPidHypotheses[i]);
    w[i] = PID->Get<COMET::IRealDatum>(name.c_str())->GetValue();
    std::cout << "   - " << name << ": " <<  w[i] << std::endl;
  }
}