/*!
* @file STools_gain.cpp
* @author Simon Vercaemer, on behalf of the SoLid collaboration.
* @date 24 Nov 2016
*/

#include "STools.h"

#include <TCanvas.h>
#include <TF1.h>
#include <TH2D.h>

#include <iostream>
#include <utility>

// defenition of PeakFindFunction in STools.h
float residualsMethodGainFinder(TH1D *ht, PeakFindFunction PeakFinder){
  char buffer[511]; // buffer to name stuff
  sprintf(buffer,"gain_Res_%s", ht->GetName());
  TH1D *hin = (TH1D*) ht->Clone(buffer);

  // do an initial rebinning to get enough stats per bin,
  // fits with less than 200 entries in the largest bin never work
  float minMaxBC = 2e2;
  if (hin->GetBinContent(hin->GetMaximumBin()) < minMaxBC){
    // Not enough statistics to do propper fitting
    std::cout<<"{Error:] Not enough statistics to do gain ";
    std::cout<<"calculation using the residuals method in ";
    std::cout<<"histogram "<<ht->GetName()<<std::endl;
    return -2;
  }

  // do peak finding
  std::vector<float> posX, posY;
  int np = PeakFinder(hin, posX, posY);
  if (np == 0) return -1;

  char fName[120];
  std::vector<TF1*> v_f;
  if (np > 4){
    // only use the 4 largest peaks
    np = 4;
  }
  //int nb = hin->GetNbinsX();
  double rlb, rhb; // range low bound, range high bound
  for (int i = 0; i < np; i ++){
    findFitRange(hin, posX[i], rlb, rhb);
    if (rlb < 0){
      // this is the pedestall, don't use it
      posX.erase(posX.begin());
      posY.erase(posY.begin());
      findFitRange(hin, posX[i], rlb, rhb);
    }

    sprintf(fName,"gaus_%i",i);
    v_f.push_back(new TF1(fName,"gaus",rlb,rhb));
    v_f[i]->SetParameters(posY[i] ,posX[i], (rhb-rlb)/2.8);
    hin->Fit(v_f[i],"rq+");
  }

  // this is curently the best estimate
  float gain = v_f[0]->GetParameter(1);
  if (np == 1){
    // it's the only estimate
    return gain;
  }

  // vector with peak positions, later used to calculate residuals
  // peaks get added when a Gaussian fit is succesfull (the gaussian
  // mean is within the fit range and not in the range of a previous fit)
  std::vector<float> nPA(1,0);
  float mu;
  for (int i = 0; i < np; i ++){
    v_f[i]->GetRange(rlb, rhb);
    mu = v_f[i]->GetParameter(1);
    if (mu < rlb || mu > rhb){
      // this is a bad fit
      continue;
    }
    bool duplicate = false;
    for (int j = 0; j < i; j ++){
      v_f[j]->GetRange(rlb,rhb);
      if (mu > rlb && mu < rhb){
        // this peak falls in the range of
        // a previously found peak,
        duplicate = true;
        break;
      }
    }
    if (!duplicate){
      nPA.push_back(mu);
    }
  }

  // use the first peak as reference,
  // the gain is going to be in that area
  float maxPos = posX.at(0);
  sprintf(buffer,"gainDet_%s", hin->GetName());
  int nBins = 1000;
  TH1D *hgd = new TH1D(buffer, buffer, nBins,
                       0.55*maxPos, 1.45*maxPos);
  std::sort(nPA.begin(), nPA.end());
  float gg, res; // gain guess, residuals
  float minRes = CalculateResiduals(nPA,gain); // minimum residual
  for (int i = 0; i < hgd->GetNbinsX(); i ++){
    gg = hgd->GetXaxis()->GetBinCenter(i+1);
    res = CalculateResiduals(nPA,gg);
    hgd->Fill(gg,res);
    if (res < minRes){
      minRes = res;
      gain = gg;
    }
  }

  hin->Write();
  hgd->Write();

  // wrap stuff up
  delete hgd;
  return gain;
}


//==============================================================================
float FirstDerivative(TH1D *h, int b){
  float dir = 0;
  if (b > 2 && b < h->GetNbinsX()-2){
    dir = 8*(h->GetBinContent(b+1) - h->GetBinContent(b-1));
    dir += h->GetBinContent(b-2) - h->GetBinContent(b+2);
    dir /= 12;
  } else if (b > 1 && b < h->GetNbinsX()-1){
    dir = 0.5*(h->GetBinContent(b+1) - h->GetBinContent(b-1));
  } else {
    if (b == 1){
      dir = h->GetBinContent(2) - h->GetBinContent(1);
    } else if (b == h->GetNbinsX()){
      dir = h->GetBinContent(b) - h->GetBinContent(b-1);
    } else {
      // b is out of range, do nothing
    }
  }
  return dir;
}


//==============================================================================
void findFitRange(TH1D *hin, float peak,
                  double &lb, double &hb){
  // vars needed for both sides
  int gb = hin->GetXaxis()->FindBin(peak); // peak bin
  float dir;  // derivative

  // lower side
  int lbb = gb-3;   // lower bound bin
  do {
    // extend range as long as:
    // * bin within range
    // * don't get into a rising part
    lbb -= 2;
    dir = FirstDerivative(hin, lbb);
  } while (dir > 0);
  lb = hin->GetXaxis()->GetBinCenter(lbb);

  // upper side
  int hbb = gb+3;  // high bound bin, number of bins
  do {
    // extend range as long as:
    // * bin within range
    // * don't get into a rising part
    hbb += 2;
    dir = FirstDerivative(hin, hbb);
  } while (dir < 0);
  hb = hin->GetXaxis()->GetBinCenter(hbb);
}

//==============================================================================
float CalculateResiduals(std::vector<float> v, float t){
  float res = 0; // residuals
  float n, nr;   // n peaks, n peaks rounded
  unsigned int ne = v.size();
  for (uint i = 0; i < ne; i ++){
    for (uint j = 0; j < ne; j ++){
      n = (v[j] - v[i])/t;
      nr = round(n);
      res += (n-nr)*(n-nr);
    }
  }
  return res;
}

//==============================================================================
bool vecPairSort(std::pair<float,float> A, std::pair<float, float> B){
  return A.second > B.second;
}

//==============================================================================
int findGainPeaks_LocMax(TH1D *hin, std::vector<float> &posX,
                            std::vector<float> &posY){
  // vector for peak positions(first) and amplitudes(second)
  std::vector< std::pair<float, float> > v_p;
  double bc, FRM; // bin content, forward range maximum
  int FRMp; // forward range maximum position
  int R = 5, WFL = hin->GetNbinsX(); // range, waveform length
  bool LM; // local maximum
  int nPeaks = 0;
  for (int i = R+1; i < WFL - R; i ++){
    bc = hin->GetBinContent(i);
    if (bc == 0) continue;
    FRM = bc;
    FRMp = i;
    LM = true;
    for (int j = 1; j <= R; j ++){
      if (hin->GetBinContent(i+j) > FRM){
        // skip bins that certainly won't be a local maximum
        FRM = hin->GetBinContent(i+j);
        FRMp = i+j;
        LM = false;
      }
      if (hin->GetBinContent(i-j) > bc) LM = false;
    }
    if (LM){
      v_p.push_back(std::make_pair(hin->GetXaxis()->GetBinCenter(i), bc));
      i += R;
      nPeaks ++;
    } else if (FRMp > i) i = FRMp - 1;
  }

  // sort peaks according to amplitude
  std::sort(v_p.begin(), v_p.end(), vecPairSort);

  // fill the position vectors
  posX.clear();
  posY.clear();
  for (int i = 0; i < nPeaks; i ++){
    posX.push_back(v_p.at(i).first);
    posY.push_back(v_p.at(i).second);
  }

  v_p.clear();
  return nPeaks;
}


//==============================================================================
double histoFoldingGainFinder(TH1D * h, bool amp){
  // Range to search for gains - problems encountered when made too high.
  // Integer steps used, since thats the precision available.
  int gainLow = 10, gainHigh = 70;
  if (!amp){
    gainLow = 50;
    gainHigh = 700;
  }

  // Create a th2 for a scan over gain trials, and another for its RMS profile.
  std::string titleMod = "hMod_" + std::string(h->GetName());
  std::string titleModRMS = "hModRMS_" + std::string(h->GetName());
  TH1D hRes(titleModRMS.c_str(), titleModRMS.c_str(),
            (gainHigh - gainLow), gainLow, gainHigh);
  TH2D hMod(titleMod.c_str(), titleMod.c_str(),
            (gainHigh - gainLow), gainLow, gainHigh, 100, -1, 1);

  // Scan over each trial gain and each bin of the input histo.
  int iGainBin = 1, nBins = h->GetXaxis()->GetNbins();
  for (float iGain = gainLow; iGain<gainHigh; iGain += 1, iGainBin++){
    for (int iBin=1; iBin<nBins; iBin++) {
      double binContent = h->GetBinContent(iBin);

      // Take this bin content, mod the bin position.
      double binPos = h->GetXaxis()->GetBinCenter(iBin)/double(iGain);
      while (binPos > 0.5) binPos -= 1.;
      hMod.Fill(iGain, binPos, binContent);
    }

    // Get the RMS of the modded input histo this gain trial.
    TH1D * hProj = hMod.ProjectionY("temp", iGainBin, iGainBin);
    double RMS = hProj->GetRMS();
    hRes.SetBinContent(iGainBin, RMS);
    delete hProj;
  }

  hMod.Write();
  hRes.Write();

  // Take the gain as the minimum bin position - an obvious place to improve.
  double gain = hRes.GetXaxis()->GetBinCenter(hRes.GetMinimumBin());

  return gain;
}