#include <iostream>
#include <string>
#include <TMath.h>
#include <TEnv.h>
#include <TSystem.h>
#include <TString.h>
#include "IoaAnalysisUtils.hxx"

ClassImp(COMET::IoaAnalysisUtils);

COMET::IoaAnalysisUtils::IoaAnalysisUtils(){}
COMET::IoaAnalysisUtils::~IoaAnalysisUtils(){}

COMET::IoaAnalysisUtils::ESubdetector COMET::IoaAnalysisUtils::PathToSubdetector(const std::string path) {
 
  if( path.find("Basket") != std::string::npos ) {
    //P0D
    if(path.find("P0D") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kP0D;
    }
    //DsECal
    else if(path.find("DsECal") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kDsECal;
    }
    //FGD1
    else if(path.find("FGD1") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kFGD1;
    }
    //FGD2
    else if(path.find("FGD2") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kFGD2;
    }
    //TPC1
    else if(path.find("TPC1") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kTPC1;
    }
    //TPC2
    else if(path.find("TPC2") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kTPC2;
    }
    //TPC3
    else if(path.find("TPC3") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kTPC3;
    }
  }
  
  if(path.find("Clam")) {
    //BrlECal
    if(path.find("BrlECal") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kBrlECal;
    }
    //P0DECal
    else if(path.find("P0DECal") != std::string::npos) {
      return COMET::IoaAnalysisUtils::kP0DECal;
    }
    //MRD
    else if((path.find("MRD") != std::string::npos) || (path.find("FluxReturn") != std::string::npos)) {
      return COMET::IoaAnalysisUtils::kMRD;
    }
  }
  
  //Other parts of the off-axis
  if(path.find("OA") != std::string::npos) {
    return COMET::IoaAnalysisUtils::kOffAxis;
  }
  //INGRID
  else if(path.find("INGRID") != std::string::npos) {
    return COMET::IoaAnalysisUtils::kINGRID;
  }
  else {
    return COMET::IoaAnalysisUtils::kCavern;
  }
}


std::string COMET::IoaAnalysisUtils::DetectorName(const ESubdetector subdet) {
  switch(subdet) {
    case kFGD1: return "FGD1"; break;
    case kFGD2: return "FGD2"; break;
    case kP0D: return "P0D"; break;
    case kDsECal: return "DsECal"; break;
    case kBrlECal: return "BrlECal"; break;
    case kP0DECal: return "P0DECal"; break;
    case kTPC1: return "TPC1"; break;
    case kTPC2: return "TPC2"; break;
    case kTPC3: return "TPC3"; break;
    case kMRD: return "MRD"; break;
    case kOffAxis: return "OffAxis"; break;
    case kINGRID: return "INGRID"; break;
    case kCavern: return "Cavern"; break;
    default: return "NONE";
  }
}


COMET::IoaAnalysisUtils::EParticleCategory COMET::IoaAnalysisUtils::PDGToCategory(const Int_t pdg) {
  TParticlePDG pdgInfo(pdg);
  return PDGInfoToCategory(&pdgInfo);
}


COMET::IoaAnalysisUtils::EParticleCategory COMET::IoaAnalysisUtils::PDGInfoToCategory(const TParticlePDG *pdgInfo) {
  if(pdgInfo) {
    Int_t pdg = pdgInfo->PdgCode();
    // rounding Double_t charge/3 to closest Int_t (ROOT returns a double in units of e/3!)
    Int_t charge = TMath::Nint(pdgInfo->Charge()/3.0) ;
    std::string particleClass = pdgInfo->ParticleClass();
    // Charged Particles
    if(charge != 0) {
      // Leptons
      if(TMath::Abs(pdg) == 11 || TMath::Abs(pdg) == 13 || TMath::Abs(pdg) == 15) {
        return kChargedLepton;
      }
      // Baryons
      if(particleClass.find("Baryon") != std::string::npos) {
        return kChargedBaryon;
      } else
      // Mesons
      if(particleClass.find("Meson") != std::string::npos) {
        return kChargedMeson;
      }
      // Other Charged
      else {
        return kOtherCharged;
      }
    } else
    // Photons
    if(pdg == 22) {
      return kPhoton;
    }
    // Other Neutral
    else {
      return kOtherNeutral;
    }
    std::cout << "Problem!! PDGInfoToCategory: pdg is " << pdg
      << ", class is " << particleClass << std::endl;
    // will never happen
    return kOther;
  }
  else {
    return kOther;
  }
}


void COMET::IoaAnalysisUtils::UseCustomPDGTable() {
  gEnv->SetValue("Root.DatabasePDG", TString(gSystem->Getenv("OAANALYSISROOT")) + "/input/oaanalysis_pdg_table.txt");
}