/* Code for determining rise time in all channels, to be compiled into a 
   running programme. Compilation:
   INC=`root-config --incdir`
   LIBS=`root-config --libs`
   g++ -std=c++11 -I${INC} ${LIBS} -o measure_delay.exe delay_analysis_compiled.cpp
- see usage with arguments in main() below
*/

#include "TMath.h"
#include "TF1.h"
#include "TSpline.h"
#include "TGraphErrors.h"
#include <algorithm>
#include <string>
#include <cmath>
#include <ctime>
#include <iostream>
#include <fstream>
#include <vector>
#include <utility>
#include <cstdlib>
#include <cstdio>
#include <fstream>

using namespace std;

vector<vector<pair<double, pair<double, double>>>> snapshots(32);

/* fname = txt file, delay = relative delay between pulses in ns, maxpos = 
   position of first peak in time bins, ndelays = number of pulses in a pulse
   train, start_peak = first peak in the pulse train to include in the analysis
*/
bool analyse_tree(const char* fname, double delay = 20, int maxpos = 15, 
		  int ndelays = 10, int start_peak = 0) {
  ifstream in(fname);
  if (in.fail()) {
    cerr << "Cannot open file named " << fname << "." << endl;
    return false;
  }
  int spacing = 15; // spacing between pulses
  int m = spacing; // time window in bins
  const int len = 206; // length of a pulse train
  const int nchan = 32;
  vector<vector<vector<double>>> data(nchan);
  int d = ndelays - start_peak;
  for (int i = 0; i < nchan; ++i) {
    vector<vector<double>> v(m * d);
    for (size_t j = 0; j < v.size(); ++j) {
      vector<double> u;
      v[j] = u;
    }
    data[i] = v;
  }
  for (int i = 0; i < nchan; ++i) {
    vector<pair<double, pair<double, double>>> v(m * (ndelays - start_peak));
    snapshots[i] = v;
  }
  bool first_zero = false, gap_reached = false;
  double n = 0;
  for (int i = 0; in.good(); ++i) {
    int sample[nchan];
    int nzeroes = 0;
    for (int j = 0; j < nchan; ++j) {
      in >> sample[j];
      if (sample[j] < 30)
	++nzeroes;
    }
    if (!gap_reached && nzeroes > 20) { // dip in most channels
      for (int k = 0; k < 3; ++k) {
	nzeroes = 0;
	for (int j = 0; j < nchan; ++j) {
	  in >> sample[j];
	  if (sample[j] < 30)
	    ++nzeroes;
	}
	if (nzeroes <= 20)
	  break;
      }
      if (nzeroes > 20) { // quadruple row of zeroes
	first_zero = true;
	i = 0;
      }
      if (i == 0)
	i = -1; // outside pulse train
      else if (first_zero) {
	gap_reached = true;
	i = 0; // start of pulse train
      }
    }
    if (!gap_reached || !first_zero)
      continue;
    ++n;
    if (i == len) { // end of pulse train
      gap_reached = false;
      i = -1;
      continue;
    }
    for (int j = 0; j < nchan; ++j) {
      int k = maxpos - m/2 + 3 + start_peak * spacing;
      if (i >= k && i - k < spacing * d)
	data[j][((i - k) % m) * d + (i - k) / m].push_back(sample[j]); 
      // works if spacing == m
    }
  }
  in.close();
  if (data[0].size() == 0) { // rare occurence of no data in the pulse train, or most channels being dead
    cerr << "Corrupt data, likely due to connection problems." << endl;
    return false;
  }
  n /= len;
  for (int j = 0; j < nchan; ++j) {
    double y0[d]; // estimate of local baseline
    double x0[d]; // position of local baseline
    for (size_t i = 0; i < d; ++i) {
      y0[i] = 0;
      x0[i] = i * spacing + 1;
    }
    for (size_t i = 0; i < 3 * d; ++i) {
      y0[i % d] += TMath::Mean(data[j][i].begin(), data[j][i].end()) / 3;
    }
    TSpline3 spline("spline", x0, y0, d);
    for (int i = 0; i < data[0].size(); ++i) {
      pair<double, pair<double, double>> ppoint;
      pair<double, double> point;
      point.first = TMath::Mean(data[j][i].begin(), data[j][i].end())
	- spline.Eval((i / d - 1) + x0[i % d]);
      point.second = TMath::RMS(data[j][i].begin(), data[j][i].end()) / 
	TMath::Sqrt(n);
      ppoint.first = (i / d + 1) * 200 - (i % d) * delay;
      ppoint.second = point;
      snapshots[j][i] = ppoint;
    }
  }
  return true;
}

Double_t Fit_Function(Double_t* Coordinates, Double_t* Parameters) { 
  Double_t Tau = Parameters[0]; //Peaking time
  Double_t time_stamp = Parameters[1];  //start-time
  Double_t bl = Parameters[2];  //baseline
  Double_t A = Parameters[3];   //peak
  Double_t N = 4; //constant
  Double_t x = Coordinates[0];
  if ((x - time_stamp) < 0)
    return bl;
  else
    return bl + A * TMath::Exp(-N * ((x - time_stamp) / Tau - 1)) *
      TMath::Power((x - time_stamp) / Tau, N);
}

// help function
bool p_sort(pair<int, double> p1, pair<int, double> p2) {
  return p1.first < p2.first;
}

// Computes rise time in given channel using fits
void get_rise_time(int channel, double& rise_time, double& error, double& amp, double& amp_error, double& ts, double& ts_error) {
  vector<pair<double, pair<double, double>>>& sv = snapshots[channel];
  vector<pair<double, double>> sv_data;
  int count = 0;
  for (size_t i = 0; i < sv.size(); ++i)
    if (sv[i].second.second > 1e-10)
      ++count;
  if (count < 1) { // dead channel, no fit is carried out
    rise_time = 0;
    error = 0;
    return;
  }
  TGraphErrors* snapshot = new TGraphErrors(count); // help graph used for fit
  count = 0;
  double max = 0, maxpos = 0;
  for (size_t i = 0; i < sv.size(); ++i) {
    if (sv[i].second.second > 1e-10) {
      snapshot->SetPoint(count, sv[i].first, sv[i].second.first);
      snapshot->SetPointError(i, 0, sv[i].second.second);
      pair<double, double> p(sv[i].first, sv[i].second.first);
      if (sv[i].second.first > max) {
	maxpos = sv[i].first;
	max = sv[i].second.first;
      }
      sv_data.push_back(p);
      ++count;
    }
  }
  sort(sv_data.begin(), sv_data.end(), p_sort);
  double min_range = sv_data[0].first;
  double max_range = sv_data[sv.size() - 1].first;
  TF1* gamma_f = new TF1("gamma_f", Fit_Function, min_range, max_range, 4);
  double param[4] = {rise_time, 800, 0, max};
  gamma_f->SetParameters(param);
  gamma_f->SetParName(0, "Peak_Time");
  gamma_f->SetParName(1, "Time_stamp");
  gamma_f->SetParName(2, "Baseline");
  gamma_f->SetParName(3, "Peak");
  gamma_f->FixParameter(2, 0);
  snapshot->Fit(gamma_f, "NQ", "", min_range, max_range);
  rise_time = gamma_f->GetParameter(0);
  error = gamma_f->GetParError(0);
  ts = gamma_f->GetParameter(1);
  ts_error = gamma_f->GetParError(1);
  amp = gamma_f->GetParameter(3);
  amp_error = gamma_f->GetParError(3);
  delete gamma_f;
  delete snapshot;
}

// Write rise time vs channel number in file with given filename
// rt_guess specifies an initial guess for the rise time
void rise_time_vs_channel(char* filename, double rt_guess = 150) {
  FILE* outfile = fopen(filename, "w");
  fprintf(outfile, "%9s %15s %7s %15s %7s %15s\n", "Channel #", "Rise time (ns)", "", "Amplitude (ADC)", "", "Timestamp (ns)");
  for (int ch = 0; ch < 32; ++ch) {
    double time = rt_guess; 
    double error = 0;
    double amp = 0; 
    double amp_error = 0;
    double ts = 0; 
    double ts_error = 0;
    get_rise_time(ch, time, error, amp, amp_error, ts, ts_error);
    fprintf(outfile, "%2d%7s %6.2f +/- %4.2f %7s %6.2f +/- %4.2f %7s %6.2f +/- %4.2f\n", ch, "", time, error, "", amp, amp_error, "", ts, ts_error);
  }
}

int main(int argc, char** argv) {
  if (argc < 4) {
    printf("Not enough arguments: argc = %d\nSyntax should be:\n./measure_delay.exe chipstring config adcfilename\n", argc);
    printf("Example:\n./measure_delay.exe 100000 20mV_delay-all /home/tpc/robot/20180725/123510_sampa100000/123510_sampa100000_trorc00_link00_20mV_delay-all_adc.txt\n");
    return -1;
  }
  char* infile = argv[3];
  // find dirname (last /)
  string str(infile);
  size_t found;
  found = str.find_last_of("/");
  char* outfile = Form("%s/rise_time_sampa%s_%s.log",  str.substr(0,found).c_str(), argv[1], argv[2]);
  string config(argv[2]);
  string muon_config("4mV_delay-all");
  int maxpos, start_peak;
  double risetime; // initial guess for the rise time
  if (config == muon_config) {
    // muon setting
    maxpos = 45;
    start_peak = 2; // skips first two peaks in pulse train
    risetime = 300;
  } else {
    // tpc setting
    maxpos = 15;
    start_peak = 0;
    risetime = 150;
  }
  if (!analyse_tree(infile, 20, maxpos, 10, start_peak))
    return 1;
  rise_time_vs_channel(outfile, risetime);
  return 0;
}