#include <RAT/AtmosphericCut.hh>
#include <RAT/DBJsonLoader.hh>
#include <RAT/DB.hh>
#include <RAT/BitManip.hh>

#include <RAT/DBTable.hh>
#include <RAT/DBTextLoader.hh>
#include <RAT/DU/TrigBits.hh>
#include <RAT/Log.hh>

namespace RAT {

  void AtmosphericCut::BeginOfRun(DS::Run& run){

    fLClean = DB::Get()->GetLink("DATACLEANING",fName);
    fNhitMin = static_cast<unsigned int>(fLClean->GetI("primary_nhit_threshold"));
    fOwlNhitMax = static_cast<unsigned int>(fLClean->GetI("owl_threshold"));
    fSecondaryNhitMin = static_cast<unsigned int>(fLClean->GetI("secondary_nhit_threshold"));
    fTimeDeadWindow = fLClean->GetD("dead_time_window");
    fTimeWindow = fLClean->GetD("time_window");
    fCutTimeWindow = fLClean->GetD("cut_time_window");
    fFlatTacCut = fLClean->GetD("time_spread_cut");

    fLastHighNHitGTID = 0;
    fInWindow = false;
    fInDeadWindow = false;
    fFoundLastFile = false;
    fFollowerCount = 0;

    // Initialize various DB/Table/Column names
    fDBName = "Atmospherics.ratdb";
    fDBTableName = "ATMOSPHERICS";
    fDBGTIDColumnName = "gtid_atmospheric";
    fDBDaysColumnName = "days_atmospheric";
    fDBSecondsColumnName = "secs_atmospheric";
    fDBNanoSecondsColumnName = "nsecs_atmospheric";
    fNumberOfFollowersName = "follower_count";
    fLastDBName = "LastAtmospheric.ratdb";
    fLastDBTableName = "LAST_ATMOSPHERIC";
    fLastGTIDName = "last_gtid_atmospheric";
    fLastDayName = "last_day_atmospheric";
    fLastSecName = "last_sec_atmospheric";
    fLastNsecName = "last_nsec_atmospheric";

    // For the second pass load the RATDB tables that were created in the first
    // pass specifiying the location in time of the atmospheric event. Also,
    // check if there was an atmospheric at the end of the previous run that
    // overlaps with the beginning of this run.
    if(fProcessPass == 2){
      bool foundDBFile = false;
      try{
        DBLinkPtr LastAtmosphericDB = DB::Get()->GetLink(fLastDBTableName);
        fLastEventStore = LastAtmosphericDB->GetI(fLastGTIDName);
        fLastDayStore = LastAtmosphericDB->GetI(fLastDayName);
        fLastSecStore = LastAtmosphericDB->GetI(fLastSecName);
        fLastNSecStore = LastAtmosphericDB->GetI(fLastNsecName);
        fFoundLastFile = true;
      }
      catch(const DBNotFoundError &e){
        warn << "AtmosphericCut::BeginOfRun: Couldn't find file containing last atmospheric "
                "in the previous run. Continuing without it." << newline;
        fFoundLastFile = false;
      }
      try{
        DBLinkPtr AtmosphericDB = DB::Get()->GetLink(fDBTableName);
        fEventStore = AtmosphericDB->GetIArray(fDBGTIDColumnName);
        foundDBFile = true;
      }
      catch(const DBNotFoundError &e){
        warn << "AtmosphericCut::BeginOfRun: Couldn't access DB. "
                "Attempting to use local file.\n";
        foundDBFile = false;
      }
      if(!foundDBFile){
        try {
          std::vector<DBTable *> contents = DBJsonLoader::parse(fDBName);
          std::vector<int> fRunRange = contents[0]->GetIArray("run_range");
          if(fRunRange[0] != int(run.GetRunID())){
            // Local file has wrong run range
            throw FileError(fDBName);
          }
          fEventStore = contents[0]->GetIArray(fDBGTIDColumnName);
        }
        catch(const FileError &e){
          warn << "AtmosphericCut::BeginOfRun: "
                  "Couldn't find file of Atmospheric GTIDS "
                  "with the correct run range. "
                  "Results of this cut won't be useful\n";
        }
        catch(const DBWrongTypeError &e){
          // Will occur if there were no atmospheric events
          warn << "AtmosphericCut::BeginOfRun: Either arrays "
                  "are empty (no atmospheric events in the run) "
                  "or database table does not exist.\n";
        }
      }
    }
  }

  Processor::Result AtmosphericCut::DSEvent(DS::Run&, DS::Entry& ds){
    bool pass = true; // If any triggered event fails, they all fail
    for(size_t iEV = 0; iEV < ds.GetEVCount(); iEV++){
      if(Event(ds, ds.GetEV(iEV)) != OKTRUE){
        pass = false;
      }
    }
    return pass ? OKTRUE : OKFALSE;
  }

  // Used for comparing times between two events. Returns the time difference
  // between the arguments in nanoseconds
  double AtmosphericCut::TimeDiff(DS::UniversalTime time_a, DS::UniversalTime time_b){
    unsigned int sec_a(time_a.GetSeconds()), sec_b(time_b.GetSeconds());
    unsigned int days_a(time_a.GetDays()), days_b(time_b.GetDays());
    double nsec_a(time_a.GetNanoSeconds()), nsec_b(time_b.GetNanoSeconds());
    // 86400 seconds in a day
    return ((days_a - days_b)*86400 + (sec_a - sec_b)*1e9) + (nsec_a - nsec_b);
  }

  // User for comparing times between two events. Used when one of the universal times
  // is not available (in this case for the last atmospheric event from previous run)
  double AtmosphericCut::TimeDiff(DS::UniversalTime time, unsigned int day,
                                  unsigned int sec, unsigned int nsec){
    unsigned int current_sec = time.GetSeconds();
    unsigned int current_day = time.GetDays();
    unsigned int current_nsec = time.GetNanoSeconds();
    // 86400 seconds in a day
    return ((current_day - day)*86400 + (current_sec - sec)*1e9) + (current_nsec - nsec);
  }

  Processor::Result AtmosphericCut::Event(DS::Entry&, DS::EV& ev){

    fPassFlag = true;
    if (fProcessPass == 1){

      unsigned int nhit = ev.GetUncalPMTs().GetCount();
      unsigned int owl_nhit = ev.GetUncalPMTs().GetOWLCount();

      // check if the event is an orphan
      if(ev.GetTrigType() == 0){
        UpdateMask(ev);
        return fPassFlag ? OKFALSE : OKTRUE;
      }

      // check trigger word to make sure there was a valid physics event.
      // throw out pedestal, and ext_a triggers
      if (BitManip::TestBit(ev.GetTrigType(), DU::TrigBits::Pedestal) ||
          BitManip::TestBit(ev.GetTrigType(), DU::TrigBits::EXTASY)) {
        UpdateMask(ev);
        return fPassFlag ? OKFALSE : OKTRUE;
      }

      // A quick way to look for flat TAC events
      double tac_sum = 0;
      for(size_t i = 0; i < nhit; i++){
        double time = ev.GetUncalPMTs().GetPMT(i).GetTime();
        tac_sum+=time;
      }
      double tac_avg = tac_sum/nhit;
      double tac_sigma = 0;
      for(size_t i = 0; i < nhit; i++){
        double time = ev.GetUncalPMTs().GetPMT(i).GetTime();
        tac_sigma += (time - tac_avg)*(time - tac_avg);
      }
      tac_sigma /= nhit;
      tac_sigma = sqrt(tac_sigma);
      if(tac_sigma > fFlatTacCut){
        UpdateMask(ev);
        return fPassFlag ? OKFALSE : OKTRUE;
      }

      // Very likely a muon, don't look for atmospheric event directly afterward
      if(!fInWindow && nhit>= fNhitMin && owl_nhit >= fOwlNhitMax){
        fInDeadWindow = true;
        fMuonTime = ev.GetUniversalTime();
        UpdateMask(ev);
        return fPassFlag ? OKFALSE : OKTRUE;
      }
      // End of the dead window after a muon
      else if(fInDeadWindow &&
         TimeDiff(ev.GetUniversalTime(), fMuonTime) > fTimeDeadWindow){
        fInDeadWindow = false;
      }

      // This is the potentially tagged atmospheric event
      if (!fInDeadWindow && !fInWindow && nhit >= fNhitMin && owl_nhit < fOwlNhitMax){
        fInWindow = true;
        fLastHighNHitGTID = ev.GetGTID();
        fTimeOfLastHighNhit = ev.GetUniversalTime();
      }
      // This is the tagged michel event, used to confirm the atmospheric tag
      else if (fInWindow && nhit >= fSecondaryNhitMin &&
              TimeDiff(ev.GetUniversalTime(), fTimeOfLastHighNhit) <= fTimeWindow){
        // Store the GTID and time of the intial tagged atmospheric event
        // but don't store it for every single follower
        if(std::find(fEventStore.begin(), fEventStore.end(), fLastHighNHitGTID) == fEventStore.end()){
          fEventStore.push_back(fLastHighNHitGTID);
          fDaysTimeStore.push_back((int)fTimeOfLastHighNhit.GetDays());
          fSecsTimeStore.push_back((int)fTimeOfLastHighNhit.GetSeconds());
          fNSecsTimeStore.push_back((int)fTimeOfLastHighNhit.GetNanoSeconds());
        }
        // Count the number of follower events
        fFollowerCount += 1;
      }
      // End of the window to look for followers, reset the window and follower count
      else if (fInWindow &&
              TimeDiff(ev.GetUniversalTime(), fTimeOfLastHighNhit) > fTimeWindow){
        fInWindow = false;
        if(fFollowerCount > 0){
          fFollowerCountStore.push_back(fFollowerCount);
        }
        fFollowerCount = 0;
      }

      return fPassFlag ? OKFALSE : OKTRUE;
    }

    else {  // fProcessPass == 2
      // For pass two just need to mark all events with a GTID within
      // the time window of the atmospheric flagged in pass 1
      int gtid = ev.GetGTID();
      DS::UniversalTime time = ev.GetUniversalTime();
      // Apply the cut for all atmospherics in the current run
      for(size_t i = 0; i < fEventStore.size(); i++){
        if (gtid == fEventStore[i]){
          fTimeOfLastHighNhit = time;
          fLastHighNHitGTID = gtid;
          fPassFlag = false;
          UpdateMask(ev);
          return fPassFlag ? OKFALSE : OKTRUE;
        }
      }
      // Apply if there was an atmospheric in the previous run the cut
      // of which overlaps with the current run.
      if(fLastEventStore != -1 && fFoundLastFile){
        if(TimeDiff(time, fLastDayStore, fLastSecStore, fLastNSecStore) < fCutTimeWindow){
          fPassFlag = false;
          UpdateMask(ev);
          return fPassFlag ? OKFALSE : OKTRUE;
        }
      }
      if(fLastHighNHitGTID && TimeDiff(time, fTimeOfLastHighNhit) < fCutTimeWindow){
        fPassFlag = false;
        UpdateMask(ev);
        return fPassFlag ? OKFALSE : OKTRUE;
      }
    }

    UpdateMask(ev);
    return fPassFlag ? OKFALSE : OKTRUE;
  }

  void AtmosphericCut::EndOfRun(DS::Run& run){

    if (fProcessPass == 1){
      // Now store the results in a ratdb file
      DBTable table(fDBTableName);
      table.SetI("version", 1);
      table.SetPassNumber(-1);
      table.SetRunRange(run.GetRunID(), run.GetRunID());
      table.SetIArray(fDBGTIDColumnName,fEventStore);
      table.SetIArray(fDBDaysColumnName,fDaysTimeStore);
      table.SetIArray(fDBSecondsColumnName,fSecsTimeStore);
      table.SetIArray(fDBNanoSecondsColumnName,fNSecsTimeStore);
      table.SetIArray(fNumberOfFollowersName, fFollowerCountStore);
      table.SaveAs(fDBName);

      // Since the time cut is 20s, we need to consider
      // that it might span run boundaries
      DBTable lastTable(fLastDBTableName);
      lastTable.SetPassNumber(-1);
      lastTable.SetI("version", 1);
      lastTable.SetRunRange(run.GetRunID()+1, run.GetRunID()+1);
      // If there are any events, store the last one with a run range
      // for the next run
      if(fEventStore.size() > 0){
        lastTable.SetI(fLastGTIDName, fEventStore.back());
        lastTable.SetI(fLastDayName, fDaysTimeStore.back());
        lastTable.SetI(fLastSecName, fSecsTimeStore.back());
        lastTable.SetI(fLastNsecName, fNSecsTimeStore.back());
      }
      // Otherwise set to invalid, which indicates to the next pass
      // that there were no atmospherics identified in the prev. run
      else{
        lastTable.SetI(fLastGTIDName, -1);
        lastTable.SetI(fLastDayName, -1);
        lastTable.SetI(fLastSecName, -1);
        lastTable.SetI(fLastNsecName,  -1);
      }
      lastTable.SaveAs(fLastDBName);
    }
  }
}