#include <Randomize.hh>

#include <RAT/DB.hh>
#include <RAT/AfterPulse.hh>
#include <RAT/Log.hh>
#include <RAT/HitPMTCollection.hh>
#include <RAT/DS/MCPMT.hh>
#include <RAT/DS/MCPE.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/DU/Utility.hh>
using namespace RAT;

AfterPulse::AfterPulse()
{

}

AfterPulse::~AfterPulse()
{

}

void AfterPulse::BeginOfRun()
{
  // Select which data entry to use
  DBLinkPtr fDataType = DB::Get()->GetLink("AFTERPULSING", "data_type");
  const string dtype = fDataType->GetS("data");

  DBLinkGroup dbTable = DB::Get()->GetLinkGroup("AFTERPULSING");

  int count_dtype = 0;
  for(DBLinkGroup::iterator iAP = dbTable.begin(); iAP != dbTable.end(); iAP++){
    const string name = iAP->first;
    DBLinkPtr fAPTable = iAP->second;
    if(name == dtype){
      fAPFlag         = fAPTable->GetI("apFlag");
      fAPFraction     = fAPTable->GetD("apFraction");
      fAPTime         = fAPTable->GetDArray("apTime");
      fAPProb         = fAPTable->GetDArray("apProb");
      count_dtype += 1;
    }
  }

  if(count_dtype != 1){
    Log::Die("After-pulsing data type not found");
  }

  if(fAPFraction >= 1){
    Log::Die("Fraction of PEs that create afterpulses must be less than 100%");
  }

  if(fAPFlag == 0 || fAPFraction == 0.0){
    info << "AfterPulsing::BeginOfRun PMT After Pulsing: OFF" << newline;
  }
  else if(fAPFlag == 1){
    info << "AfterPulsing::BeginOfRun PMT After Pulsing data using: " << dtype << newline;
    info << "AfterPulsing::BeginOfRun PMT After Pulsing: ON, " << fAPFraction*100 <<
            " percent of photoelectrons will create an afterpulse" << newline;
  }
  else{
    Log::Die("After Pulsing Flag not valid");
  }

  // initialize
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
  fPMTCount = pmtInfo.GetCount();

  fTriggerWindow = DB::Get()->GetLink("DAQ")->GetD("triggergate");

  fAfterPulseTime.resize(fPMTCount);
  fFrontEndTime.resize(fPMTCount);
  countAP.resize(fPMTCount, 0);
}

// Sample the after pulse time distribution
double AfterPulse::CalculateAfterPulseTime()
{
  const double rand = G4UniformRand();
  size_t up = 1;
  for(size_t i = 1; i < fAPTime.size(); i++){
    if(rand < fAPProb[i]){
      up = i;
      i = fAPTime.size();
    }
  }
  double time = (rand - fAPProb[(up-1)])*(fAPTime[up] - fAPTime[(up-1)])/
                (fAPProb[up] - fAPProb[(up-1)]) + fAPTime[(up-1)];
  return time;
}

void AfterPulse::DetectAfterPulse(const uint64_t& event_time)
{
  for(size_t iPMT = 0; iPMT < fPMTCount; iPMT++){ // Loop over PMTs
    size_t npe = fFrontEndTime[iPMT].size();
    for(size_t iHit = 0; iHit < npe; iHit++){ // Loop over PEs
      for(size_t iAP = 0; iAP < countAP[iPMT]; iAP++){ // Loops over APs

        uint64_t front_end_time = fFrontEndTime[iPMT][iHit];
        uint64_t after_pulse_time = fAfterPulseTime[iPMT][iAP]; // time of after-pulse
        uint64_t low_window =  event_time + front_end_time - static_cast<uint64_t>(2*fTriggerWindow);
        uint64_t high_window = event_time + front_end_time + static_cast<uint64_t>(2*fTriggerWindow);

        // If the after-pulse falls near an event window, create a PE,
        // flagged as an after-pulse. Record the time, trigger processors
        // determine whether this after-pulse actually falls into an event.
        if(after_pulse_time > low_window && after_pulse_time < high_window){
          int64_t time = after_pulse_time - event_time;
          // Build and detect the PE
          DS::MCPE photoelectron;
          photoelectron.SetCreationTime(time);
          photoelectron.SetCharge(1);
          photoelectron.SetAfterPulse(true);
          HitPMTCollection::Get()->DetectPhotoelectron(iPMT, photoelectron);
          // Don't double count, erase the PE from after-pulsing vector
          fAfterPulseTime[iPMT].erase(fAfterPulseTime[iPMT].begin()+iAP);
          countAP[iPMT]-=1;
        }

        // If an event comes along well past the time of the after-pulse
        // clear the memory of the after-pulse and corresponding front-end-time.
        if(event_time + front_end_time > after_pulse_time + high_window){
          if(iHit == npe - 1){ // only remove once
            fAfterPulseTime[iPMT].erase(fAfterPulseTime[iPMT].begin()+iAP);
            countAP[iPMT]-=1; // keep track of the fact a PE was removed
          }
          if(iAP == countAP[iPMT] - 1){ // only remove once
            fFrontEndTime[iPMT].erase(fFrontEndTime[iPMT].begin()+iHit);
          }
        }

      } // end loop over APs
    } // end loop over PEs
  } // end loop over PMTs
}

void AfterPulse::GenerateAfterPulse(const int& days, const int& secs, const double& nsecs)
{
  if(fAPFlag == 0 || fAPFraction == 0.0) // After-pulsing turned off
     return;

  // universal-time of the event in nanoseconds
  uint64_t event_time = static_cast<uint64_t>(days*24*60*60*1e9) +
                        static_cast<uint64_t>(secs*1e9) +
                        static_cast<uint64_t>(nsecs+0.5);
  // Get sorted list of hits to add after-pulsing hits to
  const std::vector<DS::MCPMT*>& pmtHits = HitPMTCollection::Get()->GetSortedPMTCollection();
  for(std::vector<DS::MCPMT*>::const_iterator iTer = pmtHits.begin(); iTer != pmtHits.end(); ++iTer){
    int pmtID = (*iTer)->GetID();
    for(size_t iPE = 0; iPE < (*iTer)->GetMCPECount(); iPE++){
      double uniRand = G4UniformRand();
      uint64_t frontEndTime = (*iTer)->GetMCPE(iPE).GetFrontEndTime();
      fFrontEndTime[pmtID].push_back(frontEndTime); // front-end-time for each PE
      if(uniRand <= fAPFraction){ // Create an after-pulse
        // 64-bit unsigned ints to hold the universal-time in ns.
        // This gives ns precision, truncates the doubles, which is OK for this purpose.
        uint64_t afterPulseTime = CalculateAfterPulseTime();
        uint64_t globalAPTime = event_time + frontEndTime + afterPulseTime;
        fAfterPulseTime[pmtID].push_back(globalAPTime); // keeps track of after-pulse time for each PMT
        countAP[pmtID]+=1; // keeps track of number of after-pulse for each PMT
      }
    }
  }
  DetectAfterPulse(event_time);
}