/********************************************
 * 
 * This Bin2Root2.cpp is specially designed for CAEN CoMPASS output
 * 
 * that all boards and channels are combined together.
 * 
 * The data file name has format DataR_run_XXX_YYY.BIN
 * where XXX is run number, YYY from 1 when data size > 1 GB
 * The first file of the run is DataR_run_XXX.BIN
 * 
 * In this case, only the first .BIN file has header. 
 * 
 * Also, the timestamp is not fully sorted. 
 * 
 * Timeoffset is hardcoded.
 * 
 * Created by Ryan Tang, goluckyryan@gmail.com
 *                   last edit @ 2024-09-25
 * 
 *********************************************/

#include "BinReader.h"
#include <sys/time.h> /** struct timeval, select() */

inline unsigned int getTime_us(){
  unsigned int time_us;
  struct timeval t1;
  struct timezone tz;
  gettimeofday(&t1, &tz);
  time_us = (t1.tv_sec) * 1000 * 1000 + t1.tv_usec;
  return time_us;
}

#include "TFile.h"
#include "TTree.h"
#include "TMacro.h"

#define MAX_MULTI  100
#define debug 0 // for > 1 number of hit to debug

#include <sstream>
#include <string>
#include <vector>
#include <utility>
#include <filesystem>
#include <regex>

//^================================================
///============ time offset
std::vector<int64_t> timeOffsetList =  {0, 0, 0, 0};

int bufferSize = 5000000;

//^================================================

std::vector<std::pair<int, std::string>> findDataFiles(const std::string& patternPrefix) {
  std::vector<std::pair<int, std::string>> results;
  std::regex pattern1(patternPrefix + "_" + R"_((\d{1,3})\.BIN)_");
  std::regex pattern2(patternPrefix + R"(\.BIN)");

  for (const auto& entry : std::filesystem::directory_iterator(".")) {
    if (entry.is_regular_file()) {
      std::string filename = entry.path().filename().string();
      std::smatch match;

      // Check for pattern DataR_run_34_YYY.BIN
      if (std::regex_match(filename, match, pattern1)) {
        int yyy = std::stoi(match[1].str());
        results.emplace_back(yyy, filename);
      } 
      // Check for pattern DataR_run_34.BIN
      else if (std::regex_match(filename, pattern2)) {
        results.emplace_back(0, filename);
      }
    }
  }

  std::sort(results.begin(), results.end(), [](const std::pair<int, std::string>& a, const std::pair<int, std::string>& b) {
    return a.first < b.first;
  });

  return results;
}


std::vector<Data> hitList_A ;
std::vector<Data> hitList_B ;
BinReader * reader = nullptr;

int ReadBatch(int bufferSize){

  if( hitList_A.size() == 0 ){
    for( int i = 0 ; i < bufferSize; i++ ){
      if( reader->ReadBlock() != 1 ){
        break;
      }
      int bd = reader->data.BoardID;
      if( bd < timeOffsetList.size() ) reader->data.TimeStamp += timeOffsetList[bd];
      hitList_A.push_back(reader->data);
    }
    if( hitList_A.size() >  1){
      std::sort(hitList_A.begin(), hitList_A.end(), [](const Data& a, const Data& b) {
        return a.TimeStamp < b.TimeStamp;
      });
    }
  }else{
    hitList_A.clear();
    hitList_A = hitList_B;
  }

  //reader next batch
  hitList_B.clear();
  for( int i = 0 ; i < bufferSize; i++ ){
    if( reader->ReadBlock() != 1 ){
      break;
    }
    int bd = reader->data.BoardID;
    if( bd < timeOffsetList.size() ) reader->data.TimeStamp += timeOffsetList[bd];
    hitList_B.push_back(reader->data);
  }
  if( hitList_B.size() > 1 ){
    std::sort(hitList_B.begin(), hitList_B.end(), [](const Data& a, const Data& b) {
      return a.TimeStamp < b.TimeStamp;
    });
  }

  if( hitList_B.size() == 0 ) return 1;

  uint64_t t0_A = hitList_A.front().TimeStamp;
  uint64_t t1_A = hitList_A.back().TimeStamp;
  uint64_t t0_B = hitList_B.front().TimeStamp;

  if( t0_A >= t0_B ){
    printf("need to increase the batch size\n");
    printf("t0_A : %15lu | %zu\n", t0_A, hitList_A.size());
    printf("t0_B : %15lu | %zu\n", t0_B, hitList_B.size());
    return -1;
  }

  if( t1_A >= t0_B ){ // move the overlap part to hitsLis_A;

    ulong ID_A = 0;
    ulong ID_B = 0;
    std::vector<Data> hitTemp;

    // find the hit that is >= t0_B, save them to hitTemp
    for( size_t j = 0; j < hitList_A.size() ; j++){
      if( hitList_A[j].TimeStamp < t0_B ) continue;;
      if( ID_A == 0 ) ID_A = j;
      hitTemp.push_back(hitList_A[j]);
    }

    // remove hitList_A element that is >= t0_B
    hitList_A.erase(hitList_A.begin() + ID_A, hitList_A.end() );
    
    // find the hit that is <= t1_A, save them to hitTemp
    for( size_t j = 0; j < hitList_B.size(); j++){
      if( hitList_B[j].TimeStamp > t1_A ) {
        break;
      }
      hitTemp.push_back(hitList_B[j]);
      ID_B = j + 1;
    }

    // remove hit elements that is <= t1_A
    hitList_B.erase(hitList_B.begin(), hitList_B.begin() + ID_B  );

    // sort hitTemp
    std::sort(hitTemp.begin(), hitTemp.end(), [](const Data& a, const Data& b) {
      return a.TimeStamp < b.TimeStamp;
    });

    //push back the hitList_A
    for( size_t j = 0; j < hitTemp.size(); j++){
      hitList_A.push_back(hitTemp[j]);
    }

  }

  // printf("\n");
  // printf("hitList A : %zu\n", hitList_A.size());
  // printf("hitList B : %zu\n", hitList_B.size());

  return 1;

}


//^#############################################################
//^#############################################################
int main(int argc, char **argv) {
  
  printf("=========================================\n");
  printf("===      CoMPASS Binary to root       ===\n");
  printf("=========================================\n");  
  if (argc < 3)    {
    printf("Incorrect number of arguments:\n");
    printf("%s [timeWindow]  [RunNum] \n", argv[0]);
    printf("       timeWindow : in ns, -1 = no event building \n");   
    printf("\n\n");

    return 1;
  }

  unsigned int runStartTime = getTime_us();

  ///============= read input
  long timeWindow = atoi(argv[1]); // ns
  short runNum = atoi(argv[2]);

  ///=========== Form outFileName;
  TString outFileName = Form("run%03d_%s.root", runNum, (timeWindow < 0 ? "single" : std::to_string(timeWindow).c_str()) ); 

  ///=========== Find all BIN file with runNum 
  std::string pattern = "DataR_run_" + std::to_string(runNum);
  printf("Searching file with pattern %s....\n", pattern.c_str());

  std::vector<std::pair<int, std::string>> inFileName = findDataFiles(pattern);
  int nFile = inFileName.size();

  printf("=========================================\n"); 
  printf("       Time Window = %ld ns = %.1f us\n", timeWindow, timeWindow/1000.);
  printf("    Max multiplity = %d hits/event (hard coded)\n", MAX_MULTI);
  printf("     Out file name = %s \n", outFileName.Data());
  printf("       Buffer Size = %d\n", bufferSize);
  printf("=============================== Number of input Files : %d \n", nFile);
  for( int i = 0 ; i < nFile; i++ ){
    printf("   %3d | %s\n", i, inFileName[i].second.c_str());
  }
  printf("========================================= Scanning files\n"); 
  unsigned long long totalHitCount = 0;
  uint16_t tempHeader;
  for( int i = 0; i < nFile; i++){
    BinReader tempReader;
    if( i == 0 ) {
      tempReader.OpenFile(inFileName[i].second);
      tempHeader = tempReader.data.Header;
    }else{
      tempReader.OpenFile(inFileName[i].second, 0, true);
      tempReader.SetCustomHeader(tempHeader);
    }
    tempReader.ScanNumHit();
    totalHitCount += tempReader.GetNumHit();
  }

  printf("============================ total Num. of Hit : %llu\n", totalHitCount); 

  timeWindow *= 1000; // to ps

  printf("========================================= Initializing tree...\n"); 
  unsigned long long                evID = 0;
  unsigned int                     multi = 0;
  unsigned short           sn[MAX_MULTI] = {0}; 
  unsigned short           ch[MAX_MULTI] = {0}; 
  unsigned short            e[MAX_MULTI] = {0};
  // unsigned short           e2[MAX_MULTI] = {0};
  unsigned long long      e_t[MAX_MULTI] = {0};
  // unsigned short          e_f[MAX_MULTI] = {0};
  // unsigned short  traceLength[MAX_MULTI];
  // short trace[MAX_MULTI][MAX_TRACE_LENGTH];

  TFile * outRootFile = new TFile(outFileName, "recreate");
  TTree * tree = new TTree("tree", outFileName);
  tree->Branch("evID",           &evID, "event_ID/l"); 
  tree->Branch("multi",         &multi, "multi/i"); 
  tree->Branch("sn",                sn, "sn[multi]/s");
  tree->Branch("ch",                ch, "ch[multi]/s");
  tree->Branch("e",                  e, "e[multi]/s");
  // tree->Branch("e2",                e2, "e2[multi]/s");
  tree->Branch("e_t",              e_t, "e_t[multi]/l");  


  // record the first time stamp and last time stamp
  unsigned long long tStart = 0;
  unsigned long long tEnd = 0;
  unsigned long long hitProcessed = 0;

  printf("========================================= Read File & events building...\n"); 
  ///=========================== Read file
  reader = new BinReader();
  int fileID = 0;
  reader->OpenFile(inFileName[fileID].second);
  const uint16_t header = reader->data.Header;

  int ret = 0;
  uint64_t t0 = -1;

  std::vector<Data> event;

  do{

    if( reader->IsEndOfFile() ){
      reader->CloseFile();
      fileID++;
      if( fileID < nFile ) {
        printf("\n========= Open next file.\n");
        reader->OpenFile(inFileName[fileID].second, 0, true);
        reader->SetCustomHeader(header);
      }else{
        printf("\n========= no more file.\n");
        // break;
      }
    }
    ret = ReadBatch(bufferSize);
    if( tStart == 0 ) tStart = hitList_A[0].TimeStamp;
    tEnd = hitList_A.back().TimeStamp;

    //build event from hitList_A;
    for( size_t i = 0; i < hitList_A.size(); i++){

      if( timeWindow < 0 ){

        multi = 1;
        sn[0] = hitList_A[i].BoardID;
        ch[0] = hitList_A[i].Channel;
        e[0] = hitList_A[i].Energy;
        // e2[p] = hitList_A[i].Energy_short;
        e_t[0] = hitList_A[i].TimeStamp / 1000;
        tree->Fill();
        evID ++;

      }else{

        if( t0 == -1 ) {
          t0 = hitList_A[i].TimeStamp;
          event.push_back(hitList_A[i]);
        }

        if( hitList_A[i].TimeStamp - t0 < timeWindow ){
          event.push_back(hitList_A[i]);
        }else{
          
          //fill an event
          multi = event.size();
          if( multi > MAX_MULTI ){
            printf("\033[31m event %lld has size = %d > MAX_MULTI = %d \033[0m\n", evID, multi, MAX_MULTI);
            // for( int k = 0 ; k < multi; k++){
            //   printf("%d | %d  %2d %llu | %llu , %llu , %llu\n", k, events[k].BoardID, events[k].Channel, events[k].TimeStamp, (long int)( events[k].TimeStamp - t0), t0, timeWindow);
            // }
            multi = MAX_MULTI;
          }

          for( size_t p = 0; p < multi ; p ++ ) {
            sn[p] = event[p].BoardID;
            ch[p] = event[p].Channel;
            e[p] = event[p].Energy;
            // e2[p] = event[p].Energy_short;
            e_t[p] = event[p].TimeStamp / 1000;
          }
          tree->Fill();
          evID ++;

          // start a new event
          event.clear();
          t0 = hitList_A[i].TimeStamp;
          event.push_back(hitList_A[i]);

        }
      }
      hitProcessed ++;
      if( hitProcessed % 10000 == 0 ) printf("hit Porcessed %llu/%llu hit....%.2f%%\n\033[A\r", hitProcessed, totalHitCount, hitProcessed*100./totalHitCount);

    }

    tree->Write();
    if( fileID >= nFile ) break;
  }while(ret == 1);

  tree->Write();

  unsigned int runEndTime = getTime_us();
  double runTime = (runEndTime - runStartTime) * 1e-6;
  printf("######################################### finished.\n");
  printf(" event building time = %.2f sec = %.2f min\n", runTime, runTime/60.);
  printf("  total events built = %llu by event builder (%llu in tree)\n", evID, tree->GetEntriesFast());
  printf(" total hit processed = %llu\n", hitProcessed);
  double tDuration_sec = (tEnd - tStart) * 1e-12;
  printf("     first timestamp = %20llu ps\n", tStart);
  printf("      last timestamp = %20llu ps\n", tEnd);
  printf(" total data duration = %.2f sec = %.2f min\n", tDuration_sec, tDuration_sec/60.);
  printf("========================================> saved to %s \n", outFileName.Data());

  TMacro info;
  info.AddLine(Form("tStart= %20llu ns",tStart));
  info.AddLine(Form("  tEnd= %20llu ns",tEnd));
  info.Write("info");
  outRootFile->Close();

  delete reader;

  return 0;

}