#include "EEQuickPi0Finder.h"

#include <iostream>
#include <TString.h>
#include <TF1.h>
#include <TCanvas.h>
#include <TPad.h>
#include <TGaxis.h>
#include <TLine.h>

#include "EEezPi0nCandidate.h"
#include "EEezCluster.h"
#include "StEEmcUtil/EEfeeRaw/EEmcEventHeader.h"
#include "StEEmcUtil/EEmcSmdMap/EEmcSmdMap.h"
#include "StEEmcUtil/EEmcGeom/EEmcGeomDefs.h"

ClassImp(EEQuickPi0Finder);
ClassImp(EEPi0TreeItem); 

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 EEQuickPi0Finder::EEQuickPi0Finder( const Char_t *name )
  : TDirectory(name,"EEMC Quick Pi0 Finder") 
{
  // Class Constructor
  mEzUtils = EEezUtils::instance();
  mSmdMap  = EEmcSmdMap::instance();
  mSmdSeedThreshold = 5.0;

  //-- TTree for storing pi0 events (with header and trigger info too) --
  mPi0Tree = new TTree("mPi0Tree","EEMC Quick and Dirty Pi0 Finder"); 
  mPi0TreeItem = new EEPi0TreeItem(); 
  mPi0Tree -> Branch("pi0fit", "EEPi0TreeItem",   &mPi0TreeItem, 8000, 1); 
  mPi0Tree -> Branch("trig",   "EEstarTrig",      &mEETrig,      8000, 1);
  mPi0Tree -> Branch("head",   "EEmcEventHeader", &mEEHead,      8000, 1);
  
}

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 Int_t EEQuickPi0Finder::Find( EEezCluster *cluster )
{
  // Returns true if the cluster is likely a pi0n, i.e. it
  // has two peaks in SMD with appropriate mass, etc...
    //

  EEezTower *seed = cluster -> getSeedTower();
  mPi0TreeItem -> eCluster = cluster -> getEnergy();
  mPi0TreeItem -> eSeed    = seed    -> getEnergy(); 
  mPi0TreeItem -> sector   = seed    -> getSector();
  mPi0TreeItem -> subsector = seed   -> getSubSector();
  mPi0TreeItem -> etabin    = seed   -> getEtabin(); 

  //-- Create a new pi0 candidate based on the cluster
  EEezPi0nCandidate *candidate = new EEezPi0nCandidate();
  candidate -> addCluster( cluster );

  //-- Require 3 or 4 SMD seeds --
  Int_t nSeedU = candidate -> numberSmdSeeds(0,10.0);
  Int_t nSeedV = candidate -> numberSmdSeeds(1,10.0);
  if ( nSeedU + nSeedV < 3 && nSeedU + nSeedV > 4 ) return 0;

  cd();

  //-- Name for cloning histograms
  TString name = cluster -> getSeedTower() -> getName();
  name += "_";
  name += mEEHead -> getEventNumber();

  mProfileU = (TH1F *)candidate -> profileU() -> Clone(name+"_Uplane");
  mProfileV = (TH1F *)candidate -> profileV() -> Clone(name+"_Vplane");

  //-- Count number of seeds --
  Int_t nu = candidate -> numberSmdSeeds( 0, mSmdSeedThreshold );
  Int_t nv = candidate -> numberSmdSeeds( 1, mSmdSeedThreshold );

  //-- No pi0 if neither plane has a seed --
  if ( nu < 1 ) return 0;
  if ( nv < 1 ) return 0;


  //-- Fit to a single peak first --
  doSinglePeakFit( cluster );

  //-- Now do double-peak fitting --
  doDoublePeakFit( cluster );

  //-- Compute kinematics
  doKinematics( cluster );
  
  //-- And fill the tree
  mPi0Tree -> Fill();

  //-- Done --
  cd (".."); return 1;

}

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 void EEQuickPi0Finder::doSinglePeakFit( EEezCluster *cluster )
{
  // Perform single peak fit for initialization purposes

  EEezTower *seed = cluster -> getSeedTower();
  Int_t sec,sub,eta;
  sec = seed -> getSector();
  sub = seed -> getSubSector();
  eta = seed -> getEtabin();

  Int_t umin,umax,vmin,vmax;
  mSmdMap -> getRangeU(sec,sub,eta, umin,umax);
  mSmdMap -> getRangeV(sec,sub,eta, vmin,vmax);
  


  //-------------------------------------------------------
  //--
  //-- Fit to single peak
  //--
  Option_t *opts = "RLQ";
  TCanvas *canvas = new TCanvas("mSingleFitCanvas","Fit to Pi0 Mass",600,300);
  canvas -> Divide(2);
  canvas -> Draw();
  
  //-- Fit the SMD U and V histograms to a single shower shape
  canvas -> cd(1);
  mProfileU -> Draw();
  Float_t ymin = mProfileU -> GetMinimum();
  Float_t ymax = mProfileU -> GetMaximum();
  TF1 *fitUsingle = mEzUtils -> fitSmdSingle ( mProfileU, opts, umin, umax );
  TLine *line = new TLine((Float_t)umin,ymin,(Float_t)umin,ymax);
  line -> SetLineColor(2);
  line -> Draw();
  line = new TLine((Float_t)umax,ymin,(Float_t)umax,ymax);
  line -> SetLineColor(2);
  line -> Draw();
  //-- Integrals on either side of peak
  Float_t u1mean = fitUsingle -> GetParameter(1);
  Float_t u1sigma = fitUsingle -> GetParameter(2);
  Float_t profileUmin = mProfileU -> GetXaxis() -> GetXmin();
  Float_t profileUmax = mProfileU -> GetXaxis() -> GetXmax();
  /*
  mProfileU -> GetXaxis() -> SetRangeUser( profileUmin, u1mean - 5.0*u1sigma );  
  Float_t uSumL = mProfileU -> Integral();
  mProfileU -> GetXaxis() -> SetRangeUser( u1mean + 5.0*u1sigma, profileUmax );
  Float_t uSumR = mProfileU -> Integral();
  mProfileU -> GetXaxis() -> UnZoom();
  */
  Float_t uSumL = Integrate( mProfileU, profileUmin, u1mean - 3.0*u1sigma );
  Float_t uSumR = Integrate( mProfileU, u1mean + 3.0*u1sigma, profileUmax ); 
  mPi0TreeItem -> usuml = uSumL;
  mPi0TreeItem -> usumr = uSumR; 

 

  canvas -> cd(2);
  mProfileV -> Draw();
  TF1 *fitVsingle = mEzUtils -> fitSmdSingle ( mProfileV, opts, vmin, vmax );
  ymin = mProfileV -> GetMinimum();
  ymax = mProfileV -> GetMaximum();
  line = new TLine((Float_t)vmin,ymin,(Float_t)vmin,ymax);
  line -> SetLineColor(2);
  line -> Draw();
  line = new TLine((Float_t)vmax,ymin,(Float_t)vmax,ymax);
  line -> SetLineColor(2);
  line -> Draw();
  Float_t v1mean = fitVsingle -> GetParameter(1);
  Float_t v1sigma = fitVsingle -> GetParameter(2);
  Float_t profileVmin = mProfileV -> GetXaxis() -> GetXmin();
  Float_t profileVmax = mProfileV -> GetXaxis() -> GetXmax();
  //-- Integrals on either side of peak
  /*
  mProfileV -> GetXaxis() -> SetRangeUser( profileVmin, v1mean - 5.0*v1sigma );  
  Float_t vSumL = mProfileV -> Integral();
  mProfileV -> GetXaxis() -> SetRangeUser( v1mean + 5.0*v1sigma, profileVmax );
  Float_t vSumR = mProfileV -> Integral();
  mProfileV -> GetXaxis() -> UnZoom();
  */
  Float_t vSumL = Integrate( mProfileV, profileVmin, v1mean - 3.0*v1sigma );
  Float_t vSumR = Integrate( mProfileV, v1mean + 3.0*v1sigma, profileVmax ); 
  mPi0TreeItem -> vsuml = vSumL;
  mPi0TreeItem -> vsumr = vSumR; 


  mSingleFitCanvas = canvas;


  //-------------------------------------------------------
  //--
  //-- Range to attempt to fit residuals over based on 
  //-- left vs right integrals
  //--
  Float_t uumin, uumax;
  if ( uSumL > uSumR ) {
    uumin=mProfileU -> GetXaxis() -> GetXmin();
    uumax=u1mean - 5.0*u1sigma;
  }
  else {
    uumin=u1mean + 5.0*u1sigma;
    uumax=mProfileU -> GetXaxis() -> GetXmax();
  }

  Float_t vvmin, vvmax;
  if ( vSumL > vSumR ) {
    vvmin=mProfileV -> GetXaxis() -> GetXmin();
    vvmax=v1mean - 5.0*v1sigma;
  }
  else {
    vvmin=v1mean + 5.0*v1sigma;
    vvmax=mProfileV -> GetXaxis() -> GetXmax();
  }


  //-------------------------------------------------------
  //--
  //-- Residual fit to single peak
  //--
  canvas = new TCanvas("mResidualFitCanvas","Fit to Pi0 Mass",600,300);
  canvas -> Divide(2);
  canvas -> Draw();

  canvas->cd(1);
  TString name = mProfileU -> GetName();
  mResidualU = (TH1F *)mProfileU -> Clone( name+"_residual" );
  mResidualU -> Add ( fitUsingle, -1.0 );
  TF1 *fitUresidual = mEzUtils -> fitSmdSingle( mResidualU, opts, uumin, uumax );
  ymin = mResidualU -> GetMinimum();
  ymax = mResidualU -> GetMaximum();
  line = new TLine( uumin, ymin, uumin, ymax ); line -> SetLineColor(4); line -> Draw();
  line = new TLine( uumax, ymin, uumax, ymax ); line -> SetLineColor(4); line -> Draw();

  canvas->cd(2);
  name = mProfileV -> GetName();
  mResidualV = (TH1F *)mProfileV -> Clone( name+"_residual" );
  mResidualV -> Add ( fitVsingle, -1.0 );

  TF1 *fitVresidual = mEzUtils -> fitSmdSingle( mResidualV, opts, vvmin, vvmax );
  ymin = mResidualV -> GetMinimum();
  ymax = mResidualV -> GetMaximum();
  line = new TLine( vvmin, ymin, vvmin, ymax ); line -> SetLineColor(4); line -> Draw();
  line = new TLine( vvmax, ymin, vvmax, ymax ); line -> SetLineColor(4); line -> Draw();

  std::cout << "mProfileU integral: " << std::endl;
  std::cout << "histogram: " << mProfileU -> Integral() << std::endl;
  std::cout << "fit:       " << fitUsingle -> GetParameter(0) << std::endl;
  std::cout << "fit resid: " << fitUresidual -> GetParameter(0) << std::endl;
  
  std::cout << "mProfileV integral: " << std::endl;
  std::cout << "histogram: " << mProfileV -> Integral() << std::endl;
  std::cout << "fit:       " << fitVsingle -> GetParameter(0) << std::endl;
  std::cout << "fit resid: " << fitVresidual -> GetParameter(0) << std::endl;

  mUinit[0].yield  = fitUsingle -> GetParameter(0);
  mUinit[0].mean   = fitUsingle -> GetParameter(1);
  mUinit[0].sigma  = fitUsingle -> GetParameter(2);
  mUinit[0].ryield = fitUsingle -> GetParameter(3);
  mUinit[0].rsigma = fitUsingle -> GetParameter(4);
  mUinit[0].min    = umin;
  mUinit[0].max    = umax;

  mPi0TreeItem -> umin1 = umin;
  mPi0TreeItem -> umax1 = umax;

  mVinit[0].yield  = fitVsingle -> GetParameter(0);
  mVinit[0].mean   = fitVsingle -> GetParameter(1);
  mVinit[0].sigma  = fitVsingle -> GetParameter(2);
  mVinit[0].ryield = fitVsingle -> GetParameter(3);
  mVinit[0].rsigma = fitVsingle -> GetParameter(4);
  mVinit[0].min    = vmin;
  mVinit[0].max    = vmax;

  mPi0TreeItem -> vmin1 = vmin;
  mPi0TreeItem -> vmax1 = vmax;

  mUinit[1].yield  = fitUresidual -> GetParameter(0);
  mUinit[1].mean   = fitUresidual -> GetParameter(1);
  mUinit[1].sigma  = fitUresidual -> GetParameter(2);
  mUinit[1].ryield = fitUresidual -> GetParameter(3);
  mUinit[1].rsigma = fitUresidual -> GetParameter(4);
  mUinit[1].min    = uumin;
  mUinit[1].max    = uumax;

  mPi0TreeItem -> umin2 = uumin;
  mPi0TreeItem -> umax2 = uumax;
  
  mVinit[1].yield  = fitVresidual -> GetParameter(0);
  mVinit[1].mean   = fitVresidual -> GetParameter(1);
  mVinit[1].sigma  = fitVresidual -> GetParameter(2);
  mVinit[1].ryield = fitVresidual -> GetParameter(3);
  mVinit[1].rsigma = fitVresidual -> GetParameter(4);
  mVinit[1].min    = vvmin;
  mVinit[1].max    = vvmax; 

  mPi0TreeItem -> vmin2 = vvmin;
  mPi0TreeItem -> vmax2 = vvmax;
  
}


/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 void EEQuickPi0Finder::doDoublePeakFit( EEezCluster *cluster )
{
    // Perform double-peak fit 
    mDoubleFitCanvas = new TCanvas("mDoubleFitCanvas","EEMC pi0 fit",600,300);
    mDoubleFitCanvas -> Divide(2);

    //-- U first -- 
    
    TString name = mProfileU -> GetName(); 
    TH1F *uhisto = (TH1F *)mProfileU -> Clone(name+"_dblefit");  
    TF1 *ufit = new TF1("eeDoublePeak",eeDoublePeak,0,288,10);
    ufit -> SetLineColor(2);
    ufit -> SetFillStyle(0); 
    ufit -> FixParameter(0, mUinit[0].yield );
    ufit -> FixParameter(1, mUinit[0].mean );
    ufit -> FixParameter(2, mUinit[0].sigma );
    ufit -> FixParameter(3, mUinit[0].ryield );
    ufit -> FixParameter(4, mUinit[0].rsigma );
    ufit -> FixParameter(5, mUinit[1].yield );
    ufit -> SetParameter(6, mUinit[1].mean);
    ufit -> FixParameter(7, mUinit[1].sigma );
    ufit -> FixParameter(8, mUinit[1].ryield );
    ufit -> FixParameter(9, mUinit[1].rsigma ); 
    
    mDoubleFitCanvas -> cd(1); 
    ufit -> SetParLimits( 6, mUinit[1].min, mUinit[1].max ); 
    
    uhisto -> Fit ( ufit, "RLQ" ); 


    //-- V next --  
    name = mProfileV -> GetName(); 
    TH1F *vhisto = (TH1F *)mProfileV -> Clone(name+"_dblefit");  
    TF1 *vfit = new TF1("eeDoublePeak",eeDoublePeak,0,288,10);
    vfit -> SetLineColor(2);
    vfit -> SetFillStyle(0); 
    vfit -> FixParameter(0, mVinit[0].yield );
    vfit -> FixParameter(1, mVinit[0].mean );
    vfit -> FixParameter(2, mVinit[0].sigma );
    vfit -> FixParameter(3, mVinit[0].ryield );
    vfit -> FixParameter(4, mVinit[0].rsigma );
    vfit -> FixParameter(5, mVinit[1].yield );
    vfit -> SetParameter(6, mVinit[1].mean);
    vfit -> FixParameter(7, mVinit[1].sigma );
    vfit -> FixParameter(8, mVinit[1].ryield );
    vfit -> FixParameter(9, mVinit[1].rsigma ); 
    
    mDoubleFitCanvas -> cd(2); 
    vfit -> SetParLimits( 6, mVinit[1].min, mVinit[1].max ); 
    
    vhisto -> Fit ( vfit, "RLQ" ); 


    //-- Save in TTree -------------------------------------
    //
    mPi0TreeItem -> u1yield  = ufit -> GetParameter(0); 
    mPi0TreeItem -> u1mean   = ufit -> GetParameter(1); 
    mPi0TreeItem -> u1sigma  = ufit -> GetParameter(2); 
    mPi0TreeItem -> u1ryield = ufit -> GetParameter(3);
    mPi0TreeItem -> u1rsigma = ufit -> GetParameter(4); 

    mPi0TreeItem -> u2yield  = ufit -> GetParameter(5); 
    mPi0TreeItem -> u2mean   = ufit -> GetParameter(6); 
    mPi0TreeItem -> u2sigma  = ufit -> GetParameter(7); 
    mPi0TreeItem -> u2ryield = ufit -> GetParameter(8);
    mPi0TreeItem -> u2rsigma = ufit -> GetParameter(9); 

    mPi0TreeItem -> uchi2   = ufit -> GetChisquare(); 
    mPi0TreeItem -> undf    = ufit -> GetNDF(); 

    mPi0TreeItem -> v1yield  = vfit -> GetParameter(0); 
    mPi0TreeItem -> v1mean   = vfit -> GetParameter(1); 
    mPi0TreeItem -> v1sigma  = vfit -> GetParameter(2); 
    mPi0TreeItem -> v1ryield = vfit -> GetParameter(3);
    mPi0TreeItem -> v1rsigma = vfit -> GetParameter(4); 

    mPi0TreeItem -> v2yield  = vfit -> GetParameter(5); 
    mPi0TreeItem -> v2mean   = vfit -> GetParameter(6); 
    mPi0TreeItem -> v2sigma  = vfit -> GetParameter(7); 
    mPi0TreeItem -> v2ryield = vfit -> GetParameter(8);
    mPi0TreeItem -> v2rsigma = vfit -> GetParameter(9); 

    mPi0TreeItem -> vchi2   = vfit -> GetChisquare(); 
    mPi0TreeItem -> vndf    = vfit -> GetNDF(); 

    mPi0TreeItem -> uprofile = uhisto; 
    mPi0TreeItem -> vprofile = vhisto; 

}

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 void EEQuickPi0Finder::doKinematics(EEezCluster *cluster)
{
  // Calculate pi0 kinematics
  
  //-- Zero out the kinematics entries in the tree item

  mPi0TreeItem -> zggu = 0.;
  mPi0TreeItem -> dggu = 0.;
  mPi0TreeItem -> zggv = 0.;
  mPi0TreeItem -> dggv = 0.;

  mPi0TreeItem -> zgg   = -1.;
  mPi0TreeItem -> dgg   = -1.;
  mPi0TreeItem -> phigg = -1.;
  mPi0TreeItem -> mass  = -1.;
  mPi0TreeItem -> xf    = -1.;
  mPi0TreeItem -> pt    = -1.;
  mPi0TreeItem -> eta   = -1.;
  mPi0TreeItem -> phi   = -1.;

  //-- Energy sharing --
  Float_t udiff = mPi0TreeItem -> u1yield - mPi0TreeItem -> u2yield;
  Float_t usum  = mPi0TreeItem -> u1yield + mPi0TreeItem -> u2yield;
  if ( usum > 0. )
    mPi0TreeItem -> zggu = TMath::Abs(udiff)/usum;

  Float_t vdiff = mPi0TreeItem -> v1yield - mPi0TreeItem -> v2yield;
  Float_t vsum  = mPi0TreeItem -> v1yield + mPi0TreeItem -> v2yield;
  if ( vsum > 0. )
    mPi0TreeItem -> zggv = TMath::Abs(vdiff)/vsum;

  Float_t zgg = ( mPi0TreeItem ->zggu + mPi0TreeItem ->zggv ) / 2.0;
  mPi0TreeItem -> zgg = zgg;

  //-- Distance between photons --

  Float_t dggu = mPi0TreeItem -> u1mean - mPi0TreeItem -> u2mean;
  Float_t dggv = mPi0TreeItem -> v1mean - mPi0TreeItem -> v2mean;
  Float_t dgg  = TMath::Sqrt( dggu*dggu + dggv*dggv );
  //-- Divide by 2 to get into cm
  dggu /= 2.;
  dggv /= 2.;
  dgg  /= 2.;
  mPi0TreeItem -> dggu = dggu;
  mPi0TreeItem -> dggv = dggv;
  mPi0TreeItem -> dgg  = dgg;
  
  //-- Opening angle between the photons
  Float_t phigg = dgg / kEEmcZSMD;
  mPi0TreeItem -> phigg = dgg / kEEmcZSMD;

  //-- NOTE: This is for 200 GeV! --
  Float_t eCluster = cluster -> getEnergy();
  mPi0TreeItem -> xf = eCluster / 200.0;
  
  Float_t mass2 = 0.5 * eCluster * eCluster * ( 1.0 - zgg ) * ( 1.0 + zgg ) *
    TMath::Sin( phigg / 2.0 ) * TMath::Sin( phigg / 2.0 );
  Float_t mass = TMath::Sqrt(mass2);

  mPi0TreeItem -> mass = mass;
}

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////

 Float_t EEQuickPi0Finder::Integrate( TH1F *h, Float_t min, Float_t max )
{
    // A method to integrate histogram contents between min and 
    // max... assumes bin width = 1.0.  Eliminates single-strip
    // (i.e. stale pedestal) bins. 

    Int_t minBin = h -> FindBin(min);
    Int_t maxBin = h -> FindBin(max);
    Float_t sum = 0.; 
    for ( Int_t i = minBin; i <= maxBin; i++ ) {
	Float_t y = h -> GetBinContent(i);
	if ( h->GetBinContent(i+1) > 0. && h->GetBinContent(i-1) > 0. ) sum += y;
    }
    return sum; 
}

/////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////