// Author: Andrew M. Rogers, NSCL 07/01/2006
// Original Design: Daniel Bazin, NSCL
//* Copyright (C) 2006-2008 Andrew M. Rogers

#include "TS800CrdcCalc.h"
#include "TS800Crdc.h"
#include "nrutil.h"

#include <math.h>
#include <iostream>

#include <TCanvas.h>
#include <TH1S.h>
#include <TRandom.h>
#include <TROOT.h>
#include <TSpline.h>

ClassImp(TS800CrdcCalc);

//////////////////////////////////////////////////////////////////////////
//                                                                      //
// TS800CrdcCalc                                                        //
//                                                                      //
//                                                                      //
//////////////////////////////////////////////////////////////////////////

extern short gNumrecError, gNumrecMessage;


//______________________________________________________________________________
TS800CrdcCalc::TS800CrdcCalc(const TS800CrdcCalc &calc) : TObject(calc)
{
  // -- Copy Constructor.
  ((TS800CrdcCalc&)calc).Copy(*this);
}


//______________________________________________________________________________
void TS800CrdcCalc::Copy(TObject &calc) const
{
  // -- Copy this method.
  
  TObject::Copy((TObject&)calc);
  ((TS800CrdcCalc&)calc).InterpSnap[0] = InterpSnap[0];
  ((TS800CrdcCalc&)calc).InterpSnap[1] = InterpSnap[1];

  ((TS800CrdcCalc&)calc).fInterpGaus    = fInterpGaus;
  ((TS800CrdcCalc&)calc).fInterpPoly    = fInterpPoly;
  ((TS800CrdcCalc&)calc).fInterpSnap    = fInterpSnap;
  ((TS800CrdcCalc&)calc).fInterpSpline3 = fInterpSpline3;
  ((TS800CrdcCalc&)calc).fSatFlat       = fSatFlat;
  ((TS800CrdcCalc&)calc).gravity_width  = gravity_width;
  ((TS800CrdcCalc&)calc).fit_width      = fit_width;
  ((TS800CrdcCalc&)calc).method         = method;
  ((TS800CrdcCalc&)calc).badpads        = badpads;
  for(Int_t i=0; i<TS800_FP_CRDC_CHANNELS; i++) ((TS800CrdcCalc&)calc).badList[i] = badList[i];
  ((TS800CrdcCalc&)calc).saturationFlat = saturationFlat;
  ((TS800CrdcCalc&)calc).name = name;

  //  ((TS800CrdcCalc&)calc).pad    = &((TS800CrdcPads&)pad;

  //  Clear();
}


//______________________________________________________________________________
void TS800CrdcCalc::InitClass(TString iname) 
{
  // -- Initials the defaults for the TS800CrdcCalc class.
  //

  name = iname;

  gravity_width  = 20;
  fit_width      = 10;
  method         = 1;
  badpads        = 0;
  saturationFlat = 9999999;

  ClearBadList();

  InterpSnap[0]  = -1;
  InterpSnap[1]  = 223;

  fInterpGaus    = kFALSE;
  fInterpPoly    = kFALSE;
  fInterpSpline3 = kFALSE;
  fInterpSnap    = kFALSE;
  fSatFlat       = kTRUE;
}


//______________________________________________________________________________
void TS800CrdcCalc::InitTree(TTree *tree) 
{
  // -- Initialize the event data TTree and set the branch addresses.
  //

  fChain = tree;
  fCurrent = -1;
  //  TNamed *aState = (TNamed*)fChain->GetUserInfo()->FindObject("Analysis State");

  b_maxpad    = fChain->GetBranch(name+".maxpad");
  b_padmax    = fChain->GetBranch(name+".padmax");
  b_padsum    = fChain->GetBranch(name+".padsum");
  b_padDist   = fChain->GetBranch(name+".padDist");
  b_x_gravity = fChain->GetBranch(name+".x_gravity");
  b_x_fit     = fChain->GetBranch(name+".x_fit");
  b_x_chi2    = fChain->GetBranch(name+".x_chi2");
}


//______________________________________________________________________________
void TS800CrdcCalc::Clear(Option_t*) 
{
  // -- Clear the relavent TS800CrdcCalc data members.
  //

  maxpad    = sqrt(-1.0);
  padmax    = -1;
  padsum    = sqrt(-1.0);
  padDist   = 225;
  x_gravity = sqrt(-1.0);
  x_fit     = sqrt(-1.0);
  x_chi2    = sqrt(-1.0);
}


//______________________________________________________________________________
void TS800CrdcCalc::GetEntry(int i) 
{
  // -- Get the ith entry for all TS800CrdcCalc branches.
  //

  b_maxpad->GetEntry(i);
  b_padmax->GetEntry(i);
  b_padsum->GetEntry(i);
  b_padDist->GetEntry(i);
  b_x_gravity->GetEntry(i);
  b_x_fit->GetEntry(i);
  b_x_chi2->GetEntry(i);
}


//______________________________________________________________________________
void TS800CrdcCalc::CalculateGravity() 
{
  // -- Calculate the positon on the CRDC using the center of gravity method.
  // There are some different algorithms that interpolate for badpads.
  // We only calculate the center of gravity for bads within the specified gravity_width.
  // padDist excludes badpads.

  Double_t calInitial[fCon->parent->channels];
  Double_t calInterpGaus[fCon->parent->channels];
  Double_t calInterpPoly[fCon->parent->channels];
  Double_t calInterpSpline3[fCon->parent->channels];
  Double_t calInterpFnc[fCon->parent->channels];

  for(Int_t c=0; c<fCon->parent->channels; c++){ calInitial[c] = pad->cal[c];}

  memset(calInterpGaus,   '\0',sizeof(calInterpGaus));
  memset(calInterpPoly,   '\0',sizeof(calInterpPoly));
  memset(calInterpSpline3,'\0',sizeof(calInterpSpline3));
  memset(calInterpFnc,    '\0',sizeof(calInterpFnc));

  maxpad = 0;

  // Reject bad pads stored in the badpad array.
  for(Int_t c=0; c<fCon->parent->channels; c++) {if(badList[c]==1) pad->cal[c] = sqrt(-1.0);}


  // Find pad at avalanche site (2 non-zero adjacent pads required)
  // If an adjacent pad is a badpad require that the next pad has data as long as the next pad is not bad.
  // 
  for (Int_t c=1; c<fCon->parent->channels-1; c++) {
    if (!isnan(pad->cal[c]) && !isnan(pad->cal[c-1]) && !isnan(pad->cal[c+1]) && pad->cal[c]>maxpad) {
      maxpad = pad->cal[c];
      padmax = c;
    }else if(!isnan(pad->cal[c])   && badList[c-1]==1 && !isnan(pad->cal[c-2]) && 
	     !isnan(pad->cal[c+1]) && pad->cal[c]>maxpad){
      maxpad = pad->cal[c];
      padmax = c;
    }else if(!isnan(pad->cal[c])   && badList[c+1]==1 && !isnan(pad->cal[c+2]) && 
	     !isnan(pad->cal[c-1]) && pad->cal[c]>maxpad){
      maxpad = pad->cal[c];
      padmax = c;
    }else if(!isnan(pad->cal[c])   && 
	     badList[c-1]==1 && badList[c-2]==1 && !isnan(pad->cal[c-3]) && 
	     ((badList[c+1]==1 && !isnan(pad->cal[c+2])) || (badList[c+1]==0 && !isnan(pad->cal[c+1])))
	     && pad->cal[c]>maxpad){
      maxpad = pad->cal[c];
      padmax = c;
    }else if(!isnan(pad->cal[c])   && 
	     badList[c+1]==1 && badList[c+2]==1 && !isnan(pad->cal[c+3]) && 
	     ((badList[c-1]==1 && !isnan(pad->cal[c-2])) || (badList[c-1]==0 && !isnan(pad->cal[c-1])))
	     && pad->cal[c]>maxpad){
      maxpad = pad->cal[c];
      padmax = c;
    }else if(!isnan(pad->cal[c])   && 
	     badList[c-1]==1 && !isnan(pad->cal[c-2]) && 
	     badList[c+1]==1 && !isnan(pad->cal[c+2]) && 
	     pad->cal[c]>maxpad){
      maxpad = pad->cal[c];
      padmax = c;
    }
    #ifdef EXPERIMENT_02023
    // Make an exception for crdc1 pad=30 since it is surrounded by badpads.
    if(c==30 && strcmp(parent->GetName(),"crdc1")==0){
      if(!isnan(pad->cal[c]) && pad->cal[c]>maxpad){
	maxpad = pad->cal[c];
	padmax = c;
      }
    }
    #endif
  }


  // If we failed to find a good maximum, do not set position
  if (isnan(padmax)) return;
  // Otherwise keep calculating

  // Calculate the distance any non-zero pad is from the padmax.
  for(Int_t c=0; c<224; c++) if(!isnan(pad->cal[c])) padDist += padmax-c;

  Int_t lowpad  = (Int_t)padmax - (Int_t)gravity_width/2;
  Int_t highpad = lowpad + (Int_t)gravity_width - 1;
  if (lowpad < 0) lowpad = 0;
  if (highpad >= fCon->parent->channels) highpad = fCon->parent->channels - 1;

  // Calculate pad center of gravity.
  // Do a simple interpolation to correct for badpads if turned on.
  Double_t sum = 0, mom = 0, div=0;
  if(fInterpPoly==kFALSE && fInterpSpline3==kFALSE && fInterpGaus==kFALSE){
    for (Int_t c=lowpad; c<= highpad; c++) {
      if (!isnan(pad->cal[c])) {
	sum += pad->cal[c];
	mom += c*pad->cal[c];
      }
    }
  }else if(fInterpGaus==kTRUE){
    // Fit a gaussian to the distribution and use it to interpolate for badpads.
    // NOT IMPLEMENTED!!!
    printf("* Gaussian interpolation is not implemented.\n");
  }else if(fInterpPoly==kTRUE){
    // Use Newton form for an interpolating polynomial using divided differences.
    Short_t  Nodes = 0;
    Double_t xNew[100],yNew[100],aNew[100];
    memset(xNew,'\0',sizeof(xNew));
    memset(yNew,'\0',sizeof(yNew));
    memset(aNew,'\0',sizeof(aNew));
    
    // Find coefficents.
    for(Int_t c=lowpad; c<=highpad; c++){ 
      if(!isnan(pad->cal[c])){
	xNew[Nodes] = (Double_t)c;
	yNew[Nodes] = (Double_t)pad->cal[c];
	aNew[Nodes] = (Double_t)pad->cal[c]; 
	Nodes++; 
      }
    }
    if(Nodes>0) Nodes--;
    for(Int_t j=1; j<=Nodes; j++){
      for(Int_t k=Nodes; k>=j; k--){aNew[k] = (aNew[k]-aNew[k-1])/(xNew[k]-xNew[k-j]);}
    }

    // Evaluate function at badpads and set the value.
    for(Int_t c=lowpad; c<=highpad; c++){ 
      if(badList[c]==1 && c>xNew[0] && c<xNew[Nodes]){
	Double_t tmpNew = aNew[Nodes];
	for(Int_t i=Nodes-1; i>=0; i--){ tmpNew = tmpNew*(c-xNew[i])+aNew[i];}
	if(tmpNew<0) tmpNew = sqrt(-1.0);
	pad->cal[c] = tmpNew;
      }
    }
    // Set arrays for snapshots.
    for(Int_t c=0; c<parent->channels; c++){ calInterpPoly[c] = pad->cal[c];}
    for(Int_t c=lowpad; c<=highpad; c++){ 
      if(c>=xNew[0] && c<=xNew[Nodes]){
	Double_t tmpNew = aNew[Nodes];
	for(Int_t i=Nodes-1; i>=0; i--){ tmpNew = tmpNew*(c-xNew[i])+aNew[i];}
	calInterpFnc[c] = tmpNew;
      }
    }
  }else if(fInterpSpline3==kTRUE){
    // Use cubic splines to interpolate for badpads.
    Short_t Nodes = 0;
    Double_t xSpl[100],ySpl[100];
    memset(xSpl,'\0',sizeof(xSpl));
    memset(ySpl,'\0',sizeof(ySpl));
    
    // Find Nodes.
    xSpl[0] = lowpad-1; ySpl[0] = 0.0;
    Nodes++;
    for(Int_t c=lowpad; c<=highpad; c++){ 
      if(!isnan(pad->cal[c])){
	xSpl[Nodes] = (Double_t)c;
	ySpl[Nodes] = (Double_t)pad->cal[c];
	Nodes++; 
      }
      if(isnan(pad->cal[c]) && badList[c]==0){
	xSpl[Nodes] = (Double_t)c;
	ySpl[Nodes] = 0.0;
	Nodes++; 
      }	
    }
    xSpl[Nodes] = highpad+1; ySpl[Nodes] = 0.0;
    Nodes++;

    // Create cubic spline function.
    TSpline3 spl("S800 CRDC Spline Interpolation",xSpl,ySpl,Nodes);
    // Evaluate function at badpads and set the value.
    for(Int_t c=lowpad; c<=highpad; c++){ 
      if(badList[c]==1){
	Double_t tmpSpl = spl.Eval(c);
	if(tmpSpl<0) tmpSpl = sqrt(-1.0);
	pad->cal[c] = tmpSpl;
      }
    }
    // Set arrays for snapshots.
    for(Int_t c=0; c<parent->channels; c++){ calInterpSpline3[c] = pad->cal[c];}
    for(Int_t c=lowpad; c<=highpad; c++){    calInterpFnc[c]     = spl.Eval(c);}
  }
  
  // If value is above saturation set it to saturation value and recalculate.
  if(fSatFlat){
    sum = 0; mom =0; div =0;
    for(Int_t c=0; c<fCon->parent->channels; c++) {
      if(pad->cal[c]>saturationFlat) {
	pad->cal[c] = saturationFlat;
      }
    }
    // If not interpolating calculate centroid with available pads.
    //    if(fInterpPoly==kFALSE){
    if(maxpad>saturationFlat) maxpad = saturationFlat;
    for (Int_t c=lowpad; c<= highpad; c++) {
      if (!isnan(pad->cal[c])) {
	sum += pad->cal[c];
	mom += c*pad->cal[c];
      }
    }
  }else{
    sum = 0; mom =0; div =0;
    for (Int_t c=lowpad; c<= highpad; c++) {
      if (!isnan(pad->cal[c])) {
	sum += pad->cal[c];
	mom += c*pad->cal[c];
      }
    }
  }

  // Display snapshots of distribution randomly and for badpad events.
  if(fInterpSnap){
    if(InterpSnap[0] == -1){
      printf("Enter a pad range (0-223): ");
      cin >> InterpSnap[0] >> InterpSnap[1];
    }
    TRandom rndSnap;
    Bool_t fSnap = kFALSE;
    for(Int_t c=lowpad; c<highpad; c++){ if(badList[c]==1) fSnap=kTRUE;}
    if(padmax>=InterpSnap[0] && padmax<=InterpSnap[1]){ fSnap=kTRUE;}
    else if(rndSnap.Rndm()<0.1) fSnap=kTRUE;

    if(fSnap && padmax>=InterpSnap[0] && padmax<=InterpSnap[1]){
      Char_t cont[100];
      Int_t snapBins = highpad-lowpad;
      TCanvas *c_cgSnap = (TCanvas*)gROOT->FindObject("c_cgSnap");
      if(!c_cgSnap){
	c_cgSnap = new TCanvas("c_cgSnap","Center of Gravity Method Snapshot");
	c_cgSnap->Divide(2,2);
      }
      TH1S *h_dist0 = new TH1S("h_dist0","Raw Distribution",snapBins,lowpad,highpad);
      TH1S *h_dist1 = new TH1S("h_dist1","Polynomial Interpolated Distribution",snapBins,lowpad,highpad);
      TH1S *h_dist2 = new TH1S("h_dist2","Saturated Distribution",snapBins,lowpad,highpad);
      TH1S *h_dist3 = new TH1S("h_dist3","Interpolation Function",snapBins,lowpad,highpad);
      TH1S *h_dist4 = new TH1S("h_dist4","Cubic Spline Interpolated Distribution",snapBins,lowpad,highpad);
      for(Int_t c=lowpad; c<=highpad; c++){ 
	if(isnan(calInitial[c])){ h_dist0->Fill(c,0);
	}else{           	  h_dist0->Fill(c,calInitial[c]);
	}
	if(isnan(calInterpPoly[c])){  h_dist1->Fill(c,0);
	}else{           	  h_dist1->Fill(c,calInterpPoly[c]);
	}
	if(isnan(calInterpSpline3[c])){  h_dist4->Fill(c,0);
	}else{           	  h_dist4->Fill(c,calInterpSpline3[c]);
	}
	if(isnan(pad->cal[c])){	  h_dist2->Fill(c,0);
	}else{           	  h_dist2->Fill(c,pad->cal[c]);
	}
	h_dist3->Fill(c,calInterpFnc[c]);
      }

      c_cgSnap->cd(1);
      h_dist0->SetFillColor(kRed); h_dist0->Draw();
      c_cgSnap->cd(2);
      if(fInterpPoly==kTRUE){    h_dist1->SetFillColor(kRed); h_dist1->Draw();}
      if(fInterpSpline3==kTRUE){ h_dist4->SetFillColor(kRed); h_dist4->Draw();}
      c_cgSnap->cd(3);
      h_dist2->SetFillColor(kRed); h_dist2->Draw();
      c_cgSnap->cd(4);
      h_dist3->SetFillColor(kRed); h_dist3->Draw();
      
      c_cgSnap->Modified();
      c_cgSnap->Update();      
      
      cout << "Continue? (y/n)" ;
      cin >> cont;
      if(strcmp(cont,"n")==0 || strcmp(cont,"N")==0) {InterpSnap[0]=-1; exit(1);}
      delete h_dist0;
      delete h_dist1;
      delete h_dist2;
      delete h_dist3;
      delete h_dist4;
      //      delete c_cgSnap;
    }
  }

  padsum    = sum;
  x_gravity = mom / sum;

  if (x_gravity < 0 || x_gravity > fCon->parent->channels) {
    maxpad = sqrt(-1.0);
    padmax = 256;
    padsum = sqrt(-1.0);
    x_gravity = sqrt(-1.0);
    return;
  }
}


//______________________________________________________________________________
void fgauss(double x, double a[], double *y, double dyda[], int na) 
{
  // Gaussian with offset
  double arg, ex, fac;
  arg   = (x - a[2]) / a[3];
  ex    = exp(-arg * arg);
  fac   = a[1] * ex * 2 * arg;
  *y    = a[1] * ex + a[4];

  dyda[1] = ex;
  dyda[2] = fac / a[3];
  dyda[3] = fac * arg / a[3];
  dyda[4] = 1;
}


//______________________________________________________________________________
void TS800CrdcCalc::CalculateFit() 
{
  // --
  //

  
  // Fitting program from Numerical Recipes (Levenberg-Marquardt method)
  void mrqmin(double x[], double y[], double sig[], int ndata, double a[], int ia[],
	      int ma, double **covar, double **alpha, double *chisq,
	      void (*funcs)(double, double[], double*, double[], int),
	      double *alamda);
  float gasdev(long *idum);
  const int MA=4;
  int i, *ia, iter, itst, j, k, mfit=MA;
  double alamda, chisq, ochisq, *x, *y, *sig, *a, **covar, **alpha;
  
  maxpad = 0;

  // Reject bad pads stored in the badpad array  
  for(Int_t c=0; c<parent->channels; c++){if(badList[c]==1) pad->cal[c] = sqrt(-1.0);}

  // Find pad at avalanche site maximum (2 non-zero adjacent pads required)
  for (int c=1; c<parent->channels-1; c++) {
    //    for (int j=0; j<badpads; j++) if (c == badpad[j]) pad->cal[c] = sqrt(-1.0);
    if (!isnan(pad->cal[c]) && pad->cal[c] > maxpad) {
      if ((!isnan(pad->cal[c-1]) && pad->cal[c-1] > maxpad/2) || (!isnan(pad->cal[c+1]) && pad->cal[c+1] > maxpad/2)) {
	maxpad = pad->cal[c];
	padmax = c;
      }
    }
  }
  // give up if we couldn't find any good maximum pad
  if (isnan(padmax)) return;
  // find pad boundaries
  int lowpad  = (int)padmax - (int)fit_width/2;
  int highpad = lowpad + (int)fit_width - 1;
  if (lowpad < 0) {
    lowpad  = 0;
    highpad = (int)padmax * 2;
  }
  if (highpad >= parent->channels) {
    highpad = parent->channels - 1;
    lowpad = (int)padmax - (parent->channels - 1 - (int)padmax);
  }
  int NPT = highpad - lowpad + 1; // Maximum number of points
  ia  = iVector(1, MA);
  a   = Vector(1, MA);
  x   = Vector(1, NPT);
  y   = Vector(1, NPT);
  sig = Vector(1, NPT);
  covar = matrix(1, MA, 1, MA);
  alpha = matrix(1, MA, 1, MA);
  
  // fill up array of data
  i = 0;
  for (j=lowpad; j<=highpad; j++) {
    // reject zeros and bad pads
    if (isnan(pad->cal[j])) continue;
    i++;
    x[i] = (double)j;
    y[i] = pad->cal[j];
    sig[i] = 50.0;
    //    sig[i] = sqrt(fabs(y[i]));
  }
  // reset number of points
  int Npoints = i;
  // forget it if there are too few points
  if (Npoints < MA) {
    free_matrix(alpha, 1, MA, 1, MA);
    free_matrix(covar, 1, MA, 1, MA);
    free_Vector(x, 1, NPT);
    free_Vector(y, 1, NPT);
    free_Vector(sig, 1, NPT);
    free_Vector(a, 1, MA);
    free_iVector(ia, 1, MA);
    return;
  }
  // 	for (i=1; i<=mfit; i++) ia[i]=1;
  gNumrecError = 0;
  // guess for fit starting point
  a[1] = maxpad;
  a[2] = padmax;
  a[3] = 3.0;
  a[4] = 0;
  //	a[1] = 3e7;
  //	a[1] = -4e5;
  //	a[2] = 2e3;
  //	a[3] = -3.0;
  // initialize fit
  alamda = -1;
  mrqmin(x, y, sig, Npoints, a, ia, MA, covar, alpha, &chisq, fgauss, &alamda);
  k = 1;
  itst = 0;
  // fitting loop
  for (;;) {
    k++;
    ochisq = chisq;
    mrqmin(x, y, sig, Npoints, a, ia, MA, covar, alpha, &chisq, fgauss, &alamda);
    // if we have encountered some kind of error during the fit
    // we forget about this event and keep going
    if (gNumrecError == 1) {
      x_fit = 0;
      x_chi2 = 0;
      free_matrix(alpha, 1, MA, 1, MA);
      free_matrix(covar, 1, MA, 1, MA);
      free_Vector(x, 1, NPT);
      free_Vector(y, 1, NPT);
      free_Vector(sig, 1, NPT);
      free_Vector(a, 1, MA);
      free_iVector(ia, 1, MA);
      return;
    }
    if (chisq > ochisq) itst = 0;
    else if (fabs(ochisq - chisq) < 0.01) itst++;
    if (itst < 4) continue;
    alamda = 0;
    mrqmin(x, y, sig, Npoints, a, ia, MA, covar, alpha, &chisq, fgauss, &alamda);
    break;
  }
  // success! get the error on the fitted position out of the covariance matrix
  //	x_error = covar[2][2];
  //	static long idum=-124325412;
  //	x_random = a[2] + 10.0 * x_error * gasdev(&idum);
  free_matrix(alpha, 1, MA, 1, MA);
  free_matrix(covar, 1, MA, 1, MA);
  free_Vector(x, 1, NPT);
  free_Vector(y, 1, NPT);
  free_Vector(sig, 1, NPT);
  free_Vector(a, 1, MA);
  free_iVector(ia, 1, MA);
  /*	double y1, y2, delta;
    delta = 4*a[3]*a[3] - 12*a[4]*a[2];
    if (delta < 0) return;
    delta = sqrt(delta);
    y1 = (-2*a[3] + delta)/6/a[4];
    y2 = (-2*a[3] - delta)/6/a[4];
    if (fabs(y1-padmax) < fabs(y2-padmax)) x_fit = y1;
    else x_fit = y2;*/
  //	x_fit = -a[2] / 2 / a[3];
  x_fit = a[2];
  x_chi2 = chisq;
  if (x_fit < 0 || x_fit > parent->channels) {
    maxpad = sqrt(-1.0);
    padmax = -1;
    padsum = sqrt(-1.0);
    x_fit  = sqrt(-1.0);
    x_chi2 = sqrt(-1.0);
  }
}