/****************************************
 * Root 6 macro: tbcFit.C               *
 * Time base corrections fit            *
 * M. Staric, September 2016            *
 ****************************************/

// constants
const double syncTimeBase = 47.163878; // [ns]
const int minHits = 1000;
const double minTimeDiff = 40; // [tdc bins]
const double maxTimeDiff = 80; // [tdc bins]
const double laserOffset = -100; // [ns]

// output directory
std::string outDir;

// ntuple variables
int scrod = 0;                                        
int channel = 0;
int window = 0;
int pixel = 0;
float tdc = 0;
float tdc0 = 0;
float adc = 0;
float width = 0;
float integral = 0;
int flag = 0;

// function protopypes
void tbcFitScrod(TTree* tree, int scrodID);
bool tbcFitChannel(TTree* tree, int scrodID, int chan);
string toString(int n);
double corTime(double tdc, int window, double* tcor);


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

/**
 * Time base corrections fit
 * @param fileName input root file with ntuple of selected hits
 * @param outputDir sub-directory name for the output of results
 */
void tbcFit(std::string fileName, std::string outputDir = "tbc")
{

  outDir = outputDir;
  if(outDir != "./") gSystem->mkdir(outDir.c_str(), kTRUE);
  
  TFile* f = TFile::Open(fileName.c_str());
  TTree* tree = (TTree*) f->Get("tree");

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

  int nbins = 128;
  TH1F* Hscrod = new TH1F("Hscrod", "scrod ID's", nbins, 0, nbins);
  int nev = (int)tree->GetEntries();
  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    if(flag == 2102) Hscrod->Fill(scrod);
  }

  for(int i = 0; i < nbins; i++) {
    if(Hscrod->GetBinContent(i+1) > 0) tbcFitScrod(tree, i);
  }

  delete Hscrod;
  
}


void tbcFitScrod(TTree* tree, int scrodID)
{
  // open output file
  
  std::string fileName = outDir + "/tbcScrod" + toString(scrodID) + ".root";
  cout<< "Fitting time base corrections for SCROD " << scrodID
      << ", output file: " << fileName <<endl;
  TFile* fout = TFile::Open(fileName.c_str(), "recreate");
  if(!fout) return;
  

  // find channels with cal pulse
  
  int nbins = 128;
  TH1F* Hchan = new TH1F("channels", "channels with cal pulse", nbins, 0, nbins);
  Hchan->GetXaxis()->SetTitle("channel number");
  int nev = (int)tree->GetEntries();
  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    if(flag == 2102 and scrod == scrodID) Hchan->Fill(channel);
  }

  Hchan->Write();

  TH1F* Hsuccess = new TH1F("success", "successfuly fitted channels", nbins, 0, nbins);
  for(int i = 0; i < nbins; i++) {
    if(Hchan->GetBinContent(i+1) > minHits) {
      bool ok = tbcFitChannel(tree, scrodID, i);
      if(ok) Hsuccess->Fill(i);
    }
  }
  Hsuccess->Write();

  fout->Close();
  
}


bool tbcFitChannel(TTree* tree, int scrodID, int chan)
{
  cout<<" channel " << chan;

  // determine outlayer removal cuts
  
  std::string name = "timeDiff_ch" + toString(chan);
  std::string forWhat = "scrod " + toString(scrodID) + " channel " + toString(chan);
  std::string title = "Cal pulse time difference vs. sample for " + forWhat;
  TProfile* prof = new TProfile(name.c_str(), title.c_str(), 256, 0, 256,
				minTimeDiff, maxTimeDiff, "S");
  prof->GetXaxis()->SetTitle("sample number");
  prof->GetYaxis()->SetTitle("time difference [tdc bins]");
  
  double tdcMean[256];
  for(int i = 0; i < 256; i++) tdcMean[i] = 128;
  double tdcSigma[256];
  for(int i = 0; i < 256; i++) tdcSigma[i] = 128;

  int nev = (int)tree->GetEntries();
  for(int iter = 0; iter < 5; iter++) {
    prof->Reset();
    for(int iev = 0; iev < nev; iev++) {
      tree->GetEntry(iev);
      if(flag == 2102 and scrod == scrodID and channel == chan) {
        int sample = (int(tdc)+window*64)%256;
        double dtdc = tdc - tdc0;
        if(fabs(dtdc - tdcMean[sample]) < 3 * tdcSigma[sample]) prof->Fill(sample, dtdc);
      }
    }
    for(int i = 0; i < 256; i++) {
      tdcMean[i] = prof->GetBinContent(i+1);
      tdcSigma[i] = prof->GetBinError(i+1);
    }
  }
  prof->Write();

  cout << " (" << prof->GetEntries() << " entries)";
  
  // fill system matrix A and right side b, and save to file as 2D and 1D histograms

  const int N = 256;
  double m[N];
  double tBin = syncTimeBase / 128;
  
  TMatrixDSym A(N);
  double b[N];
  for(int k = 0; k < N; k++) b[k] = 0;

  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    if(flag == 2102 and scrod == scrodID and channel == chan) {
      int sample = (int(tdc)+window*64)%256;
      double dtdc = tdc - tdc0;
      if(fabs(dtdc - tdcMean[sample]) > 3 * tdcSigma[sample]) continue;
      double sig = tdcSigma[sample] / tdcMean[sample];
      if(sig <= 0) continue;

      for(int k = 0; k < N; k++) m[k] = 0;
      double t1 = (tdc0 + 64 * window);
      double t2 = (tdc + 64 * window);
      int i1 = int(t1);
      m[i1%N] = 1.0 - (t1 - i1);
      int i2 = int(t2);
      m[i2%N] = t2 - i2;
      for(int k = i1 + 1; k < i2; k++) m[k%N] = 1;

      double sig2 = sig * sig;
    
      for(int j = 0; j < N; j++) {
        for(int k = 0; k < N; k++) {
          A(j,k) += m[j] * m[k] / sig2;
        }
      }
      for(int k = 0; k < N; k++) b[k] += m[k] / sig2;
    }
  }

  name = "matA_ch" + toString(chan);
  title = "Matrix for " + forWhat;
  TH2F* matA = new TH2F(name.c_str(), title.c_str(), N, 0, N, N, 0, N);
  for(int j = 0; j < N; j++) {
    for(int k = 0; k < N; k++) {
      matA->SetBinContent(j+1, N-k, A(j,k)); 
    }
  }
  matA->Write();

  name = "vecB_ch" + toString(chan);
  title = "Right side vector for " + forWhat;
  TH1F* vecB = new TH1F(name.c_str(), title.c_str(), N, 0, N);
  for(int k = 0; k < N; k++) {
    vecB->SetBinContent(k+1, b[k]); 
  }
  vecB->Write();

  // invert matrix and solve for dt

  double det = 0;
  A.Invert(&det);
  if(det == 0) {
    cout << " - fit failed" << endl;
    return false;
  }
  cout << " - fit successfull";

  double x[N];
  for(int k = 0; k < N; k++) {
    x[k] = 0;
    for(int j = 0; j < N; j++) {
      x[k] += A(k,j) * b[j];
    }
  }

  // calculate chi^2

  double chi2 = 0;
  int ndf = -N;
  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    if(flag == 2102 and scrod == scrodID and channel == chan) {
      int sample = (int(tdc)+window*64)%256;
      double dtdc = tdc - tdc0;
      if(fabs(dtdc - tdcMean[sample]) > 3 * tdcSigma[sample]) continue;
      double sig = tdcSigma[sample] / tdcMean[sample];
      if(sig <= 0) continue;
      
      for(int k = 0; k < N; k++) m[k] = 0;
      double t1 = (tdc0 + 64 * window);
      double t2 = (tdc + 64 * window);
      int i1 = int(t1);
      m[i1%N] = 1.0 - (t1 - i1);
      int i2 = int(t2);
      m[i2%N] = t2 - i2;
      for(int k = i1 + 1; k < i2; k++) m[k%N] = 1;
      double s = -1;
      for(int k = 0; k < N; k++) s += m[k] * x[k];
      chi2 += s * s / sig / sig;
      ndf++;
    }
  }
  cout << ", chi^2/ndf = " << chi2 << "/" << ndf << " = " <<chi2/ndf;
  
  // constrain sum of x to 2*syncTimeBase and calculate dt and tcor
  
  double sum = 0;
  for(int i = 0; i < N; i++) sum += x[i];
  if(sum == 0) {
    cout << " **error: sum=0" << endl;
    return false;
  }
  cout << endl;
  
  double scale = 2 * syncTimeBase / sum;
  double dt[N];
  for(int i = 0; i < N; i++) dt[i] = x[i] * scale;

  double tcor[N + 1];
  tcor[0] = 0;
  for(int i = 0; i < N; i++) tcor[i+1] = tcor[i] + dt[i];

  
  // save results to file as 2D and 1D histograms
  
  name = "invA_ch" + toString(chan);
  title = "Inverted matrix for " + forWhat;
  TH2F* invA = new TH2F(name.c_str(), title.c_str(), N, 0, N, N, 0, N);
  for(int j = 0; j < N; j++) {
    for(int k = 0; k < N; k++) {
      invA->SetBinContent(j+1, N-k, A(j,k)); 
    }
  }
  invA->Write();

  name = "dt_ch" + toString(chan);
  title = "Solution: sample time bins for " + forWhat;
  TH1F* vecX = new TH1F(name.c_str(), title.c_str(), N, 0, N);
  vecX->GetXaxis()->SetTitle("sample number");
  vecX->GetYaxis()->SetTitle("#Delta t [ns]");
  for(int k = 0; k < N; k++) {
    vecX->SetBinContent(k+1, dt[k]); 
    vecX->SetBinError(k+1, sqrt(A(k,k)) * scale); 
  }
  vecX->Write();

  name = "tcor_ch" + toString(chan);
  title = "Time base corrections for " + forWhat;
  TH1F* vecTcor = new TH1F(name.c_str(), title.c_str(), N + 1, 0, N + 1);
  vecTcor->GetXaxis()->SetTitle("sample number");
  vecTcor->GetYaxis()->SetTitle("t [ns]");
  for(int k = 0; k < N + 1; k++) {
    vecTcor->SetBinContent(k+1, tcor[k]); 
  }
  vecTcor->Write();

  // apply time base corrections to cal pulse and laser data and save histograms

  name = "calPulse_ch" + toString(chan);
  title = "Corrected cal pulse time vs. sample for " + forWhat;
  TH2F* calPulse = new TH2F(name.c_str(), title.c_str(),
			    N, 0, N, 200, 20, 24);
  calPulse->GetXaxis()->SetTitle("sample number");
  calPulse->GetYaxis()->SetTitle("t [ns]");
  
  name = "laser_ch" + toString(chan);
  title = "Corrected laser time vs. sample for " + forWhat;
  TH2F* laser = new TH2F(name.c_str(), title.c_str(),
			 N, 0, N, 200, laserOffset - 5, laserOffset + 5);
  laser->GetXaxis()->SetTitle("sample number");
  laser->GetYaxis()->SetTitle("t [ns]");
  
  for(int iev = 0; iev < nev; iev++) {
    tree->GetEntry(iev);
    if(scrod == scrodID and channel == chan) {
      int sample = (int(tdc)+window*64)%256;
      double time = corTime(tdc, window, tcor) - corTime(tdc0, window, tcor);
      if(flag == 2000) {
	laser->Fill(sample, time);
      } else if(flag == 2102) {
	calPulse->Fill(sample, time);
      }
    }
  }
  calPulse->Write();
  laser->Write();
  
  return true;
}


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


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

  int sampleNum = int(tdc);
  double frac = tdc - sampleNum;

  sampleNum += window * 64; // counted from window 0
  int n = sampleNum / 256;
  int k = sampleNum % 256;
  double time = n * syncTimeBase * 2 + tcor[k]; // from sample 0 window 0
  time += (tcor[k + 1] - tcor[k]) * frac; // add fraction
  
  return time;
  
}