// PMTCalib.cc
// Contact person: Gabriel Orebi Gann <gorebigann@lbl.gov> - Author
// See PMTCalib.hh for more details
//———————————————————————//

#include <RAT/PMTCalib.hh>
#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <RAT/DS/PMTSet.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
using namespace RAT;

PMTCalib::PMTCalib() : Processor("calibratePMT")
{
  fDoMC = false;
  fDoECA = fDoPCA = true; // By default do both (although for MC data PCA is turned off in the BeginOfRun)
}

void PMTCalib::SetI( const std::string& param, const int value )
{
  if( ( param != "all" ) && ( value != 0 && value != 1 ) )
    throw ParamInvalid( param, "value must be 0 or 1 i.e. false or true" );

  if( param == "mc" ) fDoMC = (value == 1);
  else if( param == "eca" ) fDoECA = (value == 1);
  else if( param == "pca" ) fDoPCA = (value == 1);
  else if( param == "all" ) PMTCalib::CalibStatus(value, fDoECA, fDoPCA, fDoMC);
  else throw ParamUnknown( param );
}

void PMTCalib::BeginOfRun( DS::Run& run )
{
  // Get the hit-cleaning stuff for cross-talk parameters
  fHC = DB::Get()->GetLink("HITCLEANING","hit_flags");
  // This allows you to turn-off crosstalk without turning off ECA
  fAddedCuts = fHC->GetSArray("default");

  DBLinkPtr pmtCalDB = DB::Get()->GetLink("PMTCAL");
  fECAMask = pmtCalDB->GetI("ECA_hitmask");
  fPCAMask = pmtCalDB->GetI("PCA_hitmask");

  fIsMC = run.GetMCFlag( );
  fEvNo = 0;
  if( fDoECA ){
    fECACal.BeginOfRun( run );
    for(size_t i = 0; i < fAddedCuts.size(); i++){
      if(fAddedCuts[i] == "crosstalk"){
        fCrossTalkCut.BeginOfRun(run);
      }
    }
  }
  if( fDoPCA ) fPCACal.BeginOfRun( run );
}

void PMTCalib::CalibStatus(int flag, bool &eca, bool &pca, bool &mc)
{
  eca = false;
  pca = false;
  mc = false;
  if( flag==9 ) mc = true;
  else{
    if( flag%10==1 ) eca = true;
    if( int((flag%100)/10)==1 ) pca = true;
  }
  return;
}

void PMTCalib::EventChecks(DS::EV& event)
{
  int prevcal = event.GetPMTCalType();
  bool doneeca, donepca, doneMC, donecal;
  PMTCalib::CalibStatus(prevcal, doneeca, donepca, doneMC);
  std::string status;
  donecal = true;
  if( doneMC ) status = "MC Truth";
  else{
    if( doneeca && donepca ) status = "ECA- & PCA-";
    else if( doneeca ) status = "ECA-";
    else if( donepca ) status = "PCA-";
    else{
      status = "Un-";
      donecal = false;
    }
    status += "Calibrated";
  }
  // Warn IF we're replacing already present PMTCal info
  //      OR writing PMTTruth to PMTCal
  // (Only on 1st event, so as not to swamp the screen)
  if(fEvNo == 0){
    if( donecal ) PMTCalib::DontDie(status, OVERWRITE, 0, 0);
    if( fDoMC ) PMTCalib::DontDie("", PMTTRUTH, 0,0);
  }
  // Check if the number of calibrated and uncalibrated
  // PMTs make sense
  int npmtuncal = event.GetUncalPMTs().GetAllCount();
  int npmtcal = event.GetCalPMTs().GetAllCount();
  // For a regular run we expect:
  // Data: NHit == NPMTUnCal
  //       NPMTCal == 0 if no calibrations yet run
  //       NPMTCal == NPMTUnCal if calib has already been done
  // MC:   NHit == NPMTUnCal (iff we ran the uncalibrator)
  //       (NPMTUnCal == 0: didn't uncalibrate!)
  //       NPMTCal == 0 if no calibrations yet run
  //       NPMTCal == NPMTUnCal if calib has already been done

  // Die, if the numbers of PMTCal / PMTUnCal / NHit are not as expected
  std::string mcString;
  if( fIsMC )  mcString = "MC";
  if( !fIsMC ) mcString = "";
  if( !donecal && npmtcal!=0 ) PMTCalib::Die(mcString, CALOBJECT, npmtcal, 0);
  if( donecal && npmtcal!=npmtuncal ) PMTCalib::Die(mcString, EVENTS, npmtcal, npmtuncal);
  // Can't do PCA without ECA
  if( fDoPCA && !fDoECA && !doneeca ) PMTCalib::Die("", PCANOECA, 0, 0);
  return;
}


Processor::Result PMTCalib::DSEvent(DS::Run& run, DS::Entry& ds)
{
  if( !fDoMC && !fDoECA && !fDoPCA ){
    warn << "PMTCalib::DSEvent: No Calibrations set, exiting." << newline;
    return OK;
  }
  if( fDoMC && fIsMC==0 ){
    warn<<"PMTCalib::DSEvent: Requested MC calibration on data, exiting."<<newline;
    RAT::Log::Die( "PMTCalib::Die: Encountered fatal error." );
  }

  const DU::PMTInfo& pmtList = DU::Utility::Get()->GetPMTInfo();

  if( fDoMC ){ // Copy MCHits to PMTCal without any calibration applied
    for( size_t iEvent = 0; iEvent < ds.GetMCEVCount(); iEvent++ ){
      double totalCharge = 0.0;
      DS::MCEV& mcEvent = ds.GetMCEV( iEvent );
      DS::EV& event = ds.GetEV( iEvent );
      PMTCalib::EventChecks(event);
      DS::CalPMTs calPMTs; // New set of calibrated PMTs
      for( size_t iMCHit = 0; iMCHit < mcEvent.GetMCHits().GetAllCount(); iMCHit++ ){
        DS::MCHit& mcHit = mcEvent.GetMCHits().GetAllPMT( iMCHit );
        // Create new PMTCal and fill it
        DS::PMTCal newHit;
        newHit.SetID(mcHit.GetID());
        newHit.SetQHS(mcHit.GetQHS());
        newHit.SetQHL(mcHit.GetQHL());
        newHit.SetQLX(mcHit.GetQLX());
        newHit.SetTime(mcHit.GetTime());
        newHit.SetCellID(mcHit.GetCellID());
        newHit.SetChanFlags(mcHit.GetChanFlags());
        calPMTs.AddPMT( newHit, pmtList.GetType( newHit.GetID() ) );
        if( newHit.GetQHS()>0. ) totalCharge += static_cast<double>( newHit.GetQHS() );
      }
      event.AddCalPMTs( calPMTs, MC );
      event.SetTotalCharge( totalCharge );
      fEvNo++;
    }
  }
  else{ // Apply the ECA and/or PCA calibrations and/or the crosstalk flag
    if( fDoECA ){// Apply the ECA calibration to the Uncalibrated PMTs
      for( size_t iEvent = 0; iEvent < ds.GetEVCount(); iEvent++ ){
        double totalCharge = 0.0;
        DS::CalPMTs calPMTs; // New set of calibrated PMTs
        DS::EV& event = ds.GetEV( iEvent );
        const DS::UncalPMTs& uncalPMTs = event.GetUncalPMTs();
        PMTCalib::EventChecks(event);
        for( size_t ipmt = 0; ipmt < uncalPMTs.GetAllCount(); ipmt++ ){
          int cell = uncalPMTs.GetAllPMT( ipmt ).GetCellID();
          int lcn = uncalPMTs.GetAllPMT( ipmt ).GetID();
          if( cell>15 )
            PMTCalib::Die("cell number", IDOOR, cell, 15);
          if( lcn>pmtList.GetCount() )
            PMTCalib::Die("PMT ID", IDOOR, lcn, pmtList.GetCount());
          DS::PMTCal newHit = fECACal.CalibratePMT( uncalPMTs.GetAllPMT( ipmt ) );
          calPMTs.AddPMT( newHit, pmtList.GetType( newHit.GetID() ) );
          if( newHit.GetQHS()>0. )
            totalCharge += static_cast<double>( newHit.GetQHS() );
        }
        event.AddCalPMTs( calPMTs, ECA );
        event.SetTotalCharge( totalCharge );
        fEvNo++;
      }
      // Now apply the crosstalk cut
      for(size_t i = 0; i < fAddedCuts.size(); i++){
        if(fAddedCuts[i] == "crosstalk"){
          fCrossTalkCut.CutCrossTalk(ds);
        }
      }
    }
    if( fDoPCA ){ // Apply the PCA calibration to ECA calibrated PMTs
      for( size_t iEvent = 0; iEvent < ds.GetEVCount(); iEvent++ ){
        double totalCharge = 0.0;
        DS::CalPMTs calPMTs; // New set of calibrated PMTs
        DS::EV& event = ds.GetEV( iEvent );
        PMTCalib::EventChecks(event);
        unsigned int calibrationType = event.GetPMTCalType();
        if( calibrationType == (ECA+PCA) && !fDoECA ){
          // Get the ECA calibrated data from the partially calibrated PMT set, if it exists
          if( event.GetPartialPMTCalTypes().size()>0 ){
            const DS::CalPMTs& previousCalPMTs = event.GetPartialCalPMTs(1);
            for( size_t ipmt = 0; ipmt < previousCalPMTs.GetAllCount(); ipmt++ ){
              DS::PMTCal newHit = fPCACal.CalibratePMT( previousCalPMTs.GetAllPMT( ipmt ) );
              calPMTs.AddPMT( newHit, pmtList.GetType( newHit.GetID() ) );
              if( newHit.GetQHS()>0. )
                totalCharge += static_cast<double>( newHit.GetQHS() );
            }
          }
          // set calibration type to ECA + PCA
          calibrationType = ECA + PCA;
        }
        else{
          // Most up to date ECA calibrated data is in CalPMTs
          const DS::CalPMTs& previousCalPMTs = event.GetCalPMTs();
          // Make absolutely sure that this is ECA-only calibrated data
          if( event.GetPMTCalType() != 1 ){
            warn<<"PMTCalib::DSEvent: trying to apply PCA calibration to non-ECA calibrated data."<<newline;
            warn<<"Something has gone wrong, dying."<<newline;
            RAT::Log::Die( "PMTCalib::Die: Encountered fatal error." );
          }
          for( size_t ipmt = 0; ipmt < previousCalPMTs.GetAllCount(); ipmt++ ){
            DS::PMTCal newHit = fPCACal.CalibratePMT( previousCalPMTs.GetAllPMT( ipmt ) );
            calPMTs.AddPMT( newHit, pmtList.GetType( newHit.GetID() ) );
            if( newHit.GetQHS()>0. )
              totalCharge += static_cast<double>( newHit.GetQHS() );
          }
          calibrationType += PCA;
        }
        event.AddCalPMTs( calPMTs, calibrationType);
        event.SetTotalCharge( totalCharge );
        fEvNo++;
      }
    }
    // Regardless of the calibration choice, fill the cleaned nhits
    for( size_t iEvent = 0; iEvent < ds.GetEVCount(); iEvent++ ){
      DS::EV& event = ds.GetEV( iEvent );
      CleanNHits( event );
    }
  }
  return OK;
}

// Handle errors
void PMTCalib::DontDie(std::string message, LiveFlagErr flagErr, int info1, int info2)
{
  warn << YELLOW << " Calibration SETBACK:" << CLR << newline;
  switch(flagErr)
    {
    case OVERWRITE:
      // PMTCal already contains information, to be overwritten
      warn<<"PMTCalib::DontDie"<<newline;
      warn<<"You are writing a new set of PMTCal,"<<newline;
      warn<<"the old set contained "<<message<<" PMT hits."<<newline;
      warn<<"Is this intentional?"<<newline;
      warn<<newline;
      break;
    case PMTTRUTH:
      // Writing PMTTruth to PMTCal
      warn<<"PMTCalib::DontDie"<<newline;
      warn<<"Be aware that you are filling PMTCal with true MC hits,"<<newline;
      warn<<"which have not been uncalibrated and recalibrated"<<newline;
      warn<<"i.e. are unphysical (not having the resultant smearing)"<<newline;
      warn<<"Is this intentional?"<<newline;
      break;
    }
}

// Handle FATAL errors
void PMTCalib::Die(std::string message, DieFlagErr flagErr, int info1, int info2)
{
  warn << RED << " Calibration PROBLEM:" << CLR << newline;
  switch(flagErr)
    {
    case CALOBJECT:
      // Data: claims no calib run, yet has PMTCal entries...
      warn<<"The flag claims this file has not yet been calibrated"<<newline;
      warn<<"And yet you have "<<info1<<" PMTCal objects for this event"<<newline;
      warn<<"Something is wrong."<<newline;
      break;
    case EVENTS:
      // NPMTCal != NPMTUnCal
      warn <<"Discrepancy in # events:"<<newline;
      warn<<"We have "<<info1<<" PMTCal but "<<info2<<" PMTUnCal in this event."<<newline;
      if(message=="")warn<<"The flag claims this data has already been calibrated"<<newline;
      if(message=="MC"){
        warn<<"The flag claims this MC should contain either the perfect `truth' PMT hit info";
        warn<<"or calibrated hit information"<<newline;
      }
      warn<<"And we expect 1:1 PMTCal:PMTUnCal"<<newline;
      warn<<"Something is wrong."<<newline;
      break;
    case PCANOECA:
      // PCA without ECA
      warn <<"You're asking me to perform the PCA calibration on";
      warn <<" ECA-uncalibrated data,"<<newline;
      warn <<"without running the ECA calibration first."<<newline;
      warn<<"This won't work, dying now."<<newline;
      break;
    case IDOOR:
      // lcn or cell ID not properly init
      warn<<"The "<<message<<" is out of range:"<<newline;
      warn<<"Value = "<<info1<<" but we expect a max of "<<info2<<newline;
      warn<<"This looks bad, dying now."<<newline;
      break;
    }

  RAT::Log::Die( "PMTCalib::Die: Encountered fatal error." );
}

void PMTCalib::CleanNHits(DS::EV& ev){

  int nhit_cleaned = 0;
  const DU::PMTCalStatus& hitstatus = DU::Utility::Get()->GetPMTCalStatus();

  // Loop over well-calibrated normal and HQE PMTs
  for(size_t iPMT = 0; iPMT < ev.GetCalPMTs().GetCount(); iPMT++){
    DS::PMTCal pmtCal = ev.GetCalPMTs().GetPMT(iPMT);
    if( hitstatus.GetHitStatus(pmtCal) == 0) // no flags raised
      nhit_cleaned++;
  }
  ev.SetNhitsCleaned(nhit_cleaned);
}