/* A cropped and more specialised version of delay_analysis.cpp, 
   where pulses of different delay are included in a pulse train in
   a single event. */

#include "TH1F.h"
#include "TMath.h"
#include "TF1.h"
#include "TGraphErrors.h"
#include "TMultiGraph.h"
#include "TLegend.h"
#include "TCanvas.h"
#include "TTree.h"
#include "TFile.h"
#include "TChain.h"
#include "TText.h"
#include "TStyle.h"
#include "TAxis.h"
#include "TProfile.h"
#include "TSpline.h"
#include <algorithm>
#include <string>
#include <cmath>
#include <ctime>
#include <iostream>
#include <fstream>
#include <vector>
#include <utility>
#include <cstdlib>
#include <cstdio>

using namespace std;

TFile* inFile;
TTree* tree;
vector<vector<pair<double, pair<double, double>>>> snapshots(32);
vector<TProfile*> pulse;

/* 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,
   zero_level = maximum ADC level considered as zero
*/
void analyse_tree(const char* fname, double delay = 20, int maxpos = 15, 
		  int ndelays = 10, int start_peak = 0, int zero_level = 40) {
  clock_t t = clock();
  ifstream in(fname);
  int spacing = 15;
  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;
    pulse.push_back(new TProfile(Form("pulse_ch%d", i), 
				 Form("Bin content in channel %d", i), len, 
				 -0.5, -0.5 + len));
  }
  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] < zero_level)
	++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] < zero_level)
	    ++nzeroes;
	}
	if (nzeroes <= 20)
	  break;
      }
      if (nzeroes > 20) { // triple 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) {
      pulse[j]->Fill(i, sample[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();
  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;
    }
  }
  cout << "Running time: " << (float) (clock() - t) / CLOCKS_PER_SEC << 
    " s." << endl;
}

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;
}

/* Draws a snapshot of the pulse using data from the specified channel. 
   show_errors: Show error bars.
   fit: make a Gamma(4) fit to the pulse, which is returned by the function.
   Otherwise a null pointer is returned.
   draw_output: display output in a TCanvas.
*/
TF1* snapshot_graph(int channel, bool& success, double rt_guess = 150, bool
		    show_errors = false, bool fit = false, bool draw_output =
		    true) {
  vector<pair<double, pair<double, double>>>& sv = snapshots[channel];
  int count = 0;
  for (size_t i = 0; i < sv.size(); ++i)
    if (sv[i].second.second > 1e-10)
      ++count;
  TGraphErrors* snapshot = new TGraphErrors(count);
  count = 0;
  vector<pair<double, double>> sv_data;
  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);
  FILE* outfile = fopen("snapshot_data.txt", "w");
  fprintf(outfile, "Time(ns) Amp(ADC)\n");
  for (size_t i = 0; i < sv_data.size(); ++i)
    fprintf(outfile, "%8.1f %8.1f\n", sv_data[i].first, sv_data[i].second);
  TCanvas* c1;
  char name[50], description[50];
  if (draw_output) {
    snapshot->SetMarkerStyle(7);
    snapshot->SetMarkerSize(2);
    snapshot->SetMarkerColor(2);
    snapshot->SetLineColor(2);
    sprintf(name, "snapsnot_canvas_ch%d", channel);
    sprintf(description, "Snapshot of pulse in channel %d", channel);
    c1 = new TCanvas(name, description, 900, 600);
  }
  TF1* gamma_f = nullptr;
  if (fit) {
    double min_range = sv_data[0].first;
    double max_range = sv_data[sv.size() - 1].first;
    gamma_f = new TF1("gamma_f", Fit_Function, min_range, max_range, 4);
    double param[4] = {rt_guess, 770 - rt_guess + 400 * (rt_guess > 200), 0,
		       max};
    gamma_f->SetParameters(param);
    gamma_f->SetParName(0, "#tau");
    gamma_f->SetParName(1, "t_{0}");
    gamma_f->SetParName(2, "b");
    gamma_f->SetParName(3, "A");
    gamma_f->FixParameter(2, 0);
    gamma_f->SetLineColor(kBlack);
    gamma_f->SetLineWidth(2);
    gamma_f->SetLineStyle(2);
    string option = draw_output ? "" : "N";
    success = snapshot->Fit("gamma_f", option.c_str(), "", min_range, max_range)
      == 0;
  }
  if (!show_errors)
    for (size_t i = 0; i < sv.size(); ++i)
      snapshot->SetPointError(i, 0, 0);
  if (draw_output) {
    snapshot->Draw("AP");
    snapshot->GetXaxis()->SetTitle("time (ns)");
    snapshot->GetYaxis()->SetTitle("amplitude (ADC)");
    //snapshot->SetTitle(description);
    snapshot->SetTitle("");
    snapshot->GetXaxis()->SetRangeUser(0, 2000);
    gStyle->SetOptFit(1);
    c1->Modified();
  }
  return gamma_f;
}

// Procedure for studying the pulse train, averaged over all events
void pulse_window(int channel) {
  TCanvas* c1 = new TCanvas(Form("c%d", channel), "Display canvas", 900, 600);
  gStyle->SetOptStat(0);
  pulse[channel]->SetTitle(Form("Display of channel %d", channel));
  pulse[channel]->Draw("bar");
}

// Draws rise time vs channel
bool rise_time_vs_channel(char* filename, double rt_guess = 150, bool save =
			  true) {
  TGraphErrors* graph = new TGraphErrors(32);
  bool success = true;
  for (int i = 0; i < 32; ++i) {
    bool fit_success;
    TF1* fit = snapshot_graph(i, fit_success, rt_guess, true, true, false);
    graph->SetPoint(i, i, fit->GetParameter(0));
    graph->SetPointError(i, 0, fit->GetParError(0));
    delete fit;
    if (!fit_success)
      success = false;
  }
  TCanvas* c2 = new TCanvas("rise_time_canvas", "Rise time vs channel #", 900,
			    600);
  graph->SetMarkerStyle(20);
  graph->Draw("AP");
  graph->GetXaxis()->SetTitle("channel #");
  graph->GetXaxis()->SetRangeUser(-1, 32);
  graph->GetYaxis()->SetTitle("rise time (ns)");
  //graph->SetTitle("Rise time vs channel #");
  graph->SetTitle("");
  c2->Modified();
  if (save)
    c2->SaveAs(filename);
  return success;
}