#include <RAT/PeakFinder.hh>
#include <RAT/Log.hh>
#include <RAT/DS/DataNotFound.hh>

#include <math.h>

#include <TH1F.h>

using namespace RAT;
using namespace std;

/// Returns the Peak value in the th1 histogram by re-calculating the mean of the
/// distribution within +/- nSigmas until convergence is achieved. The number of
/// counts and the RMS in this window are passed to PeakCounts and tRMS.
/// PeakFindWide does a broad search in the whole histogram range.
/// This is the same method used in SNO in QSNO/snomanpp/qio_tree for the rch file
/// production. It should NOT be applied to the single channel search, since, due to
/// the low stat,  it's more likely to find other peaks. In that case, a focused
/// search around t = 0 should be performed, with PeakFindNarrow, after subtracting a
/// global offset that can be obtained  with PeakFindWide.
DS::SOCPMT::EPeakFindStatus PeakFinder::PeakFindWide( const TH1F *th1, float &tPeak, float &tRMS, float &peakCount, const float nSigmas )
{

  DS::SOCPMT::EPeakFindStatus flag = DS::SOCPMT::eUnknown;

  float eps = 0.05;

  int nxbins = static_cast<int>( th1->GetNbinsX() );

  float maxnxbins = static_cast<float>( th1->GetBinLowEdge(nxbins) + th1->GetBinWidth(nxbins) );

  float   mAvg = 0, mSigma = 0;
  tPeak = 0.0; tRMS = 0.0; peakCount = 0.0;

  // Maximum number of iterations
  int maxiters = 30;

  int lowBin = 1; // ROOT first bin

  int highBin = 0, maxBin;

  bool hasConverged = false;

  float rms = 0.0, rms2;

  int nIter = 0;

  float oldPeak = 0.0;

  float fsumwx2, fsumw, fsumwx, fw;
  float fsumwx2bin, fsumwxbin, fx, fxbin;

  maxBin = nxbins;

  highBin = nxbins;

  if( th1->GetEntries() > 5 )
    {

      while( !hasConverged && nIter < maxiters )
        {

          rms2        = 0.0;
          fsumwx2     = 0.0;
          fsumw       = 0.0;
          fsumwx      = 0.0;
          fsumwxbin   = 0.0;
          fsumwx2bin  = 0.0;

          for( int i = lowBin; i < highBin + 1; i++ )
            {
              fxbin = static_cast<float>(i);
              fx = static_cast<float>( th1->GetBinCenter(i) );
              fw = static_cast<float>( th1->GetBinContent(i) );

              fsumwx2bin += fxbin * fxbin * fw;
              fsumwx2 += fx * fx * fw;
              fsumwx += fx * fw;
              fsumwxbin += fxbin * fw;
              fsumw += fw;
            }

          if( fsumw == 0 )
            {
              // No entries in histogram
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              flag = DS::SOCPMT::eNoData;
              return flag;
            }
          else
            {
              rms2 = fabs( fsumwx2bin / fsumw - fsumwxbin * fsumwxbin /( fsumw * fsumw ) );
              rms = sqrt(rms2);
              mAvg = fsumwxbin / fsumw;
              mSigma = fabs( fsumwx2 / fsumw - fsumwx * fsumwx / ( fsumw * fsumw ) );
              mSigma = sqrt( mSigma );
              tPeak = fsumwx / fsumw;
            }

          if( fabs(oldPeak - tPeak) < eps ) hasConverged = true;

          if( lowBin == highBin )
            {
              // problem with histogram boundaries
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              flag = DS::SOCPMT::eBoundaryError;
              return flag;
            }

          if( tPeak > maxnxbins )
            {
              // Peak found outside of histogram window
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              flag = DS::SOCPMT::ePeakOutsideOfWindow;
              return flag;
            }

          oldPeak = tPeak;

          lowBin  = static_cast<int>( mAvg - nSigmas * rms );
          highBin = static_cast<int>( mAvg + nSigmas * rms );

          if( lowBin < 1 ) lowBin = 1;

          if( highBin > maxBin ) highBin = maxBin;

          nIter++;

        }

      if( hasConverged )
        {
          //Calculate the number of counts in Peak +/- nSigmas;
          fsumw   = 0.0;
          for ( int j = lowBin; j < highBin; j++ ) fsumw += static_cast<float>( th1->GetBinContent(j) );

          tRMS = mSigma;
          peakCount = fsumw;

          flag = DS::SOCPMT::eSuccess;
          return flag;
        }
      else
        {
          // Did not converged after 30 iterations
          tPeak = DS::INVALID;
          tRMS  = DS::INVALID;
          peakCount = DS::INVALID;
          flag = DS::SOCPMT::eNoConvergence;
          return flag;
        }

    }
  else
    {
      // Less than 5 entries in histogram
      tPeak = DS::INVALID;
      tRMS = DS::INVALID;
      peakCount = DS::INVALID;
      flag = DS::SOCPMT::eTooFewEntries;
      return flag;
    }

}

/// Returns the Peak value in the th1 histogram by re-calculating the mean of the
/// distribution within +/- nSigmas until convergence is achieved. The number of
/// counts and the RMS in this window are passed to peakCount and tRMS.
/// PeakFindNarrow does a focused search around t = 0.
DS::SOCPMT::EPeakFindStatus PeakFinder::PeakFindNarrow(const TH1F *th1, float &tPeak, float &tRMS, float &peakCount, const float tWindow )
{

  DS::SOCPMT::EPeakFindStatus flag = DS::SOCPMT::eUnknown;

  float eps = 0.05;

  int nxbins = static_cast<int>( th1->GetNbinsX() );

  float mSigma;

  // Maximum number of iterations
  int maxiters = 30;

  int lowBin,highBin,maxBin;

  bool hasConverged = false;

  int nIter = 0;

  float oldPeak = 0.;

  tPeak = 0.;
  tRMS = 0.;
  peakCount = 0.;

  float fsumwx2, fsumw, fsumwx, fw;
  float fx;

  maxBin = nxbins;

#if ROOT_VERSION_CODE >= ROOT_VERSION(6, 0, 0)
  const TAxis *axis = th1->GetXaxis();
#else
  TAxis* axis = th1->GetXaxis();
#endif

  // Only try to find peak if histo has more than 5 entries
  if( th1->GetEntries() > 5 )
    {
      while( !hasConverged && nIter < maxiters )
        {

          fsumwx2 = 0.0;
          fsumw = 0.0;
          fsumwx = 0.0;

          lowBin = axis->FindBin(tPeak - tWindow);
          highBin = axis->FindBin(tPeak + tWindow);

          if( lowBin < 1 ) lowBin = 1; // ROOT TH1 first bin
          if( highBin > maxBin) highBin = maxBin; // ROOT TH1 last bin

          for(int i = lowBin; i < highBin + 1; i++ )
            {
              fx = static_cast<float>( th1->GetBinCenter(i) );
              fw = static_cast<float>( th1->GetBinContent(i) );

              fsumwx2 += fx * fx * fw;
              fsumwx += fx * fw;
              fsumw += fw;

            }

          if( fsumw == 0 )
            {
              // No data found around zero
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              flag = DS::SOCPMT::eNoData;
              return flag;
            }
          else
            {
              mSigma = fabs( fsumwx2 / fsumw - fsumwx * fsumwx / ( fsumw * fsumw ) );
              mSigma = sqrt( mSigma );
              tPeak   = fsumwx / fsumw;
            }

          if( fabs(oldPeak - tPeak) < eps ) hasConverged = true;

          if( lowBin == highBin )
            {
              // Problem with histo boundaries
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              flag = DS::SOCPMT::eBoundaryError;
              return flag;
            }

          oldPeak = tPeak;

          nIter++;

        }

      if( hasConverged )
        {
          tRMS = mSigma;
          peakCount = fsumw;
          flag = DS::SOCPMT::eSuccess;
          return flag;
        }
      else
        {
          // Did not converged after 30 iterations
          tPeak = DS::INVALID;
          tRMS = DS::INVALID;
          peakCount = DS::INVALID;
          flag = DS::SOCPMT::eNoConvergence;
          return flag;
        }

    }
  else
    {
      // Histogram with less than 5 entries
      tPeak = DS::INVALID;
      tRMS  = DS::INVALID;
      peakCount = DS::INVALID;
      flag = DS::SOCPMT::eTooFewEntries;
      return flag;
    }

}

/// Returns the Peak value of the SOC PMT hit times by re-calculating the mean of the
/// distribution within +/- nSigmas until convergence is achieved. The number of
/// counts and the RMS in this window are passed to peakCount and tRMS.
/// PeakFindNarrow does a focused search around t = 0.
DS::SOCPMT::EPeakFindStatus PeakFinder::PeakFindNarrow( DS::SOCPMT& socPMT, float &tPeak, float &tRMS, float &peakCount, const float globalOffset, const float tWindow )
{

  DS::SOCPMT::EPeakFindStatus flag = DS::SOCPMT::eUnknown;

  // This is the threshold for the difference between the
  // last found peak value and the current found peak value.
  // When the difference between the values of two subsequent iterations
  // is less than 'eps', the algorithm finishes.
  float eps = 0.05;

  vector< Float_t > hitTimes = socPMT.GetTimes();
  int nTimes = static_cast<int>( hitTimes.size() );
  double pmtTOF = socPMT.GetTOFManipulator();

  float mSigma;

  // Maximum number of iterations
  int maxiters = 30;

  float lowTime,highTime;

  bool hasConverged = false;

  int nIter = 0;

  float oldPeak = 0.;

  tPeak = 0.;
  tRMS = 0.;
  peakCount = 0.;

  float fsumwx2, fsumw, fsumwx;
  float fx;

  // Only try to find peak if histo has more than 5 entries
  if( nTimes > 5 )
    {
      while( !hasConverged && nIter < maxiters )
        {

          fsumwx2 = 0.0;
          fsumw = 0.0;
          fsumwx = 0.0;

          lowTime = tPeak - tWindow;
          highTime = tPeak + tWindow;

          for( int i = 0; i < nTimes; i++ )
            {
              fx = hitTimes[ i ] - pmtTOF - globalOffset;
              if ( fx > lowTime && fx < highTime ){
                fsumwx2 += fx * fx;
                fsumwx += fx;
                fsumw += 1;
              }
            }

          if( fsumw == 0 )
            {
              // No data found around zero
              tPeak = DS::INVALID;
              tRMS = DS::INVALID;
              peakCount = DS::INVALID;
              return DS::SOCPMT::eNoData;
            }
          else
            {
              mSigma = fabs( fsumwx2 / fsumw - fsumwx * fsumwx / ( fsumw * fsumw ) );
              mSigma = sqrt( mSigma );
              tPeak   = fsumwx / fsumw;
            }

          if( fabs(oldPeak - tPeak) < eps ) hasConverged = true;

          oldPeak = tPeak;

          nIter++;

        }

      if( hasConverged )
        {
          tRMS = mSigma;
          peakCount = fsumw;
          return DS::SOCPMT::eSuccess;
        }
      else
        {
          // Did not converged after 30 iterations
          tPeak = DS::INVALID;
          tRMS = DS::INVALID;
          peakCount = DS::INVALID;
          return DS::SOCPMT::eNoConvergence;
        }

    }
  else
    {
      // SOCPMT with less than 5 time entries
      tPeak = DS::INVALID;
      tRMS  = DS::INVALID;
      peakCount = DS::INVALID;
      return DS::SOCPMT::eTooFewEntries;
    }

}

/// Finds the main peak in th1 by using a sliding window (width must be
/// supplied as tWindow) over the histogram.
/// The center of the region with the most counts is returned.
/// The number of counts and RMS in the window are passed to peakCount and tRMS.
bool PeakFinder::PeakFindSlidingWindow( TH1F* const th1, float &tPeak, float &tRMS, float &peakCount, const float tWindow )
{
  RAT::info << "PeakFinder: PeakFindSlidingWindow()\n";

  if( !th1 ) return false;
  if( th1->GetEntries() <= 5 ) return false;

  int nbins = th1->GetNbinsX(), binlo, binhi;
  float startt = th1->GetBinLowEdge(1);
  float stopt = th1->GetBinLowEdge( nbins-1 ) + th1->GetBinWidth( nbins - 1);
  float t, deltat = 0.5;
  float maxSum = 0., maxTlo=0., maxRms=0.,sumwin, rmswin;

  for( t = startt; t < stopt; t+= deltat)
    {
      binlo  = th1->GetXaxis()->FindBin( t );
      binhi  = th1->GetXaxis()->FindBin( t + tWindow );

      th1->GetXaxis()->SetRange(binlo, binhi);
      sumwin = th1->Integral();
      rmswin = th1->GetRMS();

      if( sumwin > maxSum )
        {
          maxSum = sumwin;
          maxTlo = t;
          maxRms = rmswin;

        }
    }

  tRMS       = maxRms;
  peakCount = maxSum;

  tPeak = maxTlo + tWindow/2.;

  return true;

}