/**********************************************
 * Root 6 macro: rewriteNtuple.C              *
 * Include time base corrected time and width *     
 * M. Staric, September 2016                  *
 **********************************************/

// ntuple variables
int scrod = 0;      // scrodID                                      
int channel = 0;    // in: channel within scrod [0:127], out:channel within module [0:512]
int window = 0;     // ASIC window number
int pixel = 0;      // pixel ID
float tdc = 0;      // uncorrected time [tdc bins]
float tdc0 = 0;     // uncorrected time of the first cal pulse [tdc bins]
float adc = 0;      // pulse height [adc counts]
float width = 0;    // width [tdc bins]
float integral = 0; // integral [adc counts]
int flag = 0;       // flag
float tbcTime = 0;  // corrected time [ns]
float tbcWidth = 0; // corrected width [ns]
int sample = 0;     // sample number

// function prototypes
double corTime(double tdc, int window, TH1F* tcor);
std::string toString(int n);

// -- code ---------------------------------------------------------

/**
 * Rewrite ntuple: add time base corrected time and width, and sample number
 * @param inputFileName input root file
 * @param outputFileName output root file
 * @param dir path to time base correction files
 */
void rewriteNtuple(std::string inputFileName,
		   std::string outputFileName,
		   std::string dir = "tbc")
{

  if(inputFileName == outputFileName) {
    cout << "Input and output file names are the same!" << endl;
    return;
  }
  
  gROOT->CloseFiles();
  gROOT->Reset();

  // open input file and get ntuple
  
  TFile* fin = TFile::Open(inputFileName.c_str());
  if(!fin) return;
  TTree* tree = (TTree*) fin->Get("tree");
  if(!tree) return;
  tree->SetBranchAddress("scrod", &scrod);
  tree->SetBranchAddress("channel", &channel);
  tree->SetBranchAddress("window", &window);
  tree->SetBranchAddress("pixel", &pixel);
  tree->SetBranchAddress("tdc", &tdc);
  tree->SetBranchAddress("tdc0", &tdc0);
  tree->SetBranchAddress("adc", &adc);
  tree->SetBranchAddress("width", &width);
  tree->SetBranchAddress("integral", &integral);
  tree->SetBranchAddress("flag", &flag);

  // map scrod ID with boardstack number
  
  std::map<int, int> scrodToBS; // scrod ID to boardstack number
  int nev = (int)tree->GetEntries();
  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    int bs = ((pixel - 1) % 64) / 16;
    scrodToBS[scrod] = bs;
  }

  // get time base corrections (assume cal pulse in channel 7)
  
  std::vector<TH1F*> corrections(512, 0);
  typedef std::map<int, int>::iterator Iterator;
  for(Iterator it = scrodToBS.begin(); it != scrodToBS.end(); it++) {
    int scrodID = it->first;
    int bs = it->second;
    std::string fName = dir + "/tbcScrod" + toString(scrodID) + ".root";
    TFile* fcor = TFile::Open(fName.c_str());
    if(!fcor) {
      cout<< "Time base corrections for scrod " << scrodID << "not found" << endl;
      continue;
    }
    for(int k = 0; k < 16; k++) {
      int chan = k * 8 + 7;
      std::string hname = "tcor_ch" + toString(chan);
      TH1F* h = (TH1F*) fcor->Get(hname.c_str());
      if(!h) continue;
      for(int i = 0; i < 8; i++) corrections[bs * 128 + chan - i] = h;
    }
  }

  // count channels with tbc given
  
  int nn = 0;
  for(const auto& correction: corrections) {
    if(correction) nn++;
  }
  if(nn == 0) {
    cout << "No time base correctios available for any channel" << endl;
    return true;
  }
  cout << "Time base corrections available for " << nn << " channels" << endl;
  
  // open output file and book ntuple

  TFile* fout = TFile::Open(outputFileName.c_str(), "recreate");
  if(!fout) return;
  TTree* outTree = new TTree("tree", "with time base corrected time");
  outTree->Branch("scrod", &scrod, "scrod/I");
  outTree->Branch("channel", &channel, "channel/I");
  outTree->Branch("window", &window, "window/I");
  outTree->Branch("pixel", &pixel, "pixel/I");
  outTree->Branch("tdc", &tdc, "tdc/F");
  outTree->Branch("tdc0", &tdc0, "tdc0/F");
  outTree->Branch("adc", &adc, "adc/F");
  outTree->Branch("width", &width, "width/F");
  outTree->Branch("integral", &integral, "integral/F");
  outTree->Branch("flag", &flag, "flag/I");
  outTree->Branch("tbcTime", &tbcTime, "tbcTime/F");
  outTree->Branch("tbcWidth", &tbcWidth, "tbcWidth/F");
  outTree->Branch("sample", &sample, "sample/I");

  // rewite

  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    int bs = ((pixel - 1) % 64) / 16;
    channel += 128 * bs;
    auto* tcor = corrections[channel];
    if(!tcor) continue;
    tbcTime = corTime(tdc, window, tcor) - corTime(tdc0, window, tcor);
    tbcWidth = corTime(tdc + width, window, tcor) - corTime(tdc, window, tcor);
    sample = (int(tdc)+window*64)%256;
    outTree->Fill();
  }
  outTree->Write();
  fout->Close();
  fin->Close();
  
}


string toString(int n)
{
  stringstream ss;
  ss << n;
  string str;
  ss >> str;
  return str;
}

bool wait()
{
  std::cout << "Type <CR> to continue or Q to quit ";
  char qq[10];
  std::cin.getline(qq, 10);
  if (strcmp(qq, "q") == 0 || strcmp(qq, ".q") == 0 || strcmp(qq, "Q") == 0) return false;
  return true;
}

double corTime(double tdc, int window, TH1F* tcor)
{

  if(!tcor) return 0;
  
  int sampleNum = int(tdc);
  double frac = tdc - sampleNum;

  sampleNum += window * 64; // counted from window 0
  int n = sampleNum / 256;
  int k = sampleNum % 256;
  double period = tcor->GetBinContent(257); // = 2 * syncTimeBase
  double time = n * period + tcor->GetBinContent(k+1); // from sample 0 of window 0
  time += (tcor->GetBinContent(k+2) - tcor->GetBinContent(k+1)) * frac; // add fraction
  
  return time;

}