#include "EventDiscriminator.h"
#include "Event.h"

#include "TH1F.h"
#include "TFile.h"

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

EventDiscriminator::EventDiscriminator() {
}


EventDiscriminator::~EventDiscriminator() {
}


// create histograms
void EventDiscriminator::book() {
  hSig_ptSum   = new TH1F( "hSig_ptSum"  , "hSig_ptSum"  , 20, 0.0, 100.0 );
  hBkg_ptSum   = new TH1F( "hBkg_ptSum"  , "hBkg_ptSum"  , 20, 0.0, 100.0 );
  hSig_ptMin   = new TH1F( "hSig_ptMin"  , "hSig_ptMin"  , 20, 0.0, 100.0 );
  hBkg_ptMin   = new TH1F( "hBkg_ptMin"  , "hBkg_ptMin"  , 20, 0.0, 100.0 );
  hSig_ptMax   = new TH1F( "hSig_ptMax"  , "hSig_ptMax"  , 20, 0.0, 100.0 );
  hBkg_ptMax   = new TH1F( "hBkg_ptMax"  , "hBkg_ptMax"  , 20, 0.0, 100.0 );
  hSig_phiDMin = new TH1F( "hSig_phiDMin", "hSig_phiDMin", 20, 0.0, 4.0 );
  hBkg_phiDMin = new TH1F( "hBkg_phiDMin", "hBkg_phiDMin", 20, 0.0, 4.0 );
  hSig_phiDMax = new TH1F( "hSig_phiDMax", "hSig_phiDMax", 20, 0.0, 4.0 );
  hBkg_phiDMax = new TH1F( "hBkg_phiDMax", "hBkg_phiDMax", 20, 0.0, 4.0 );
  hSig_etaDMin = new TH1F( "hSig_etaDMin", "hSig_etaDMin", 20, 0.0, 5.0 );
  hBkg_etaDMin = new TH1F( "hBkg_etaDMin", "hBkg_etaDMin", 20, 0.0, 5.0 );
  hSig_etaDMax = new TH1F( "hSig_etaDMax", "hSig_etaDMax", 20, 0.0, 5.0 );
  hBkg_etaDMax = new TH1F( "hBkg_etaDMax", "hBkg_etaDMax", 20, 0.0, 5.0 );
}


// fill histograms for "signal" or "background" events
void EventDiscriminator::fill( const Event& ev, EventType type ) {

  // compute event variables
  compute( ev );
  switch ( type ) {
  case signal:
    // fill histograms for "signal" events
    hSig_ptSum  ->Fill( ptSum   );
    hSig_ptMin  ->Fill( ptMin   );
    hSig_ptMax  ->Fill( ptMax   );
    hSig_phiDMin->Fill( phiDMin );
    hSig_phiDMax->Fill( phiDMax );
    hSig_etaDMin->Fill( etaDMin );
    hSig_etaDMax->Fill( etaDMax );
    break;
  case background:
    // fill histograms for "background" events
    hBkg_ptSum  ->Fill( ptSum   );
    hBkg_ptMin  ->Fill( ptMin   );
    hBkg_ptMax  ->Fill( ptMax   );
    hBkg_phiDMin->Fill( phiDMin );
    hBkg_phiDMax->Fill( phiDMax );
    hBkg_etaDMin->Fill( etaDMin );
    hBkg_etaDMax->Fill( etaDMax );
    break;
  }
  return;

}


// save histograms
void EventDiscriminator::save( const std::string& name, EventType type ) {
  // save current working area
  TDirectory* currentDir = gDirectory;
  // open file
  TFile* file = new TFile( name.c_str(), "UPDATE" );
  switch ( type ) {
  case signal:
    // normalize histograms for "signal" events
    normalizeHisto( hSig_ptSum );
    normalizeHisto( hSig_ptMin );
    normalizeHisto( hSig_ptMax );
    normalizeHisto( hSig_phiDMin );
    normalizeHisto( hSig_phiDMax );
    normalizeHisto( hSig_etaDMin );
    normalizeHisto( hSig_etaDMax );
    // save histograms for "signal" events
    hSig_ptSum  ->Write();
    hSig_ptMin  ->Write();
    hSig_ptMax  ->Write();
    hSig_phiDMin->Write();
    hSig_phiDMax->Write();
    hSig_etaDMin->Write();
    hSig_etaDMax->Write();
    break;
  case background:
    // normalize histograms for "background" events
    normalizeHisto( hBkg_ptSum );
    normalizeHisto( hBkg_ptMin );
    normalizeHisto( hBkg_ptMax );
    normalizeHisto( hBkg_phiDMin );
    normalizeHisto( hBkg_phiDMax );
    normalizeHisto( hBkg_etaDMin );
    normalizeHisto( hBkg_etaDMax );
    // save histograms for "background" events
    hBkg_ptSum  ->Write();
    hBkg_ptMin  ->Write();
    hBkg_ptMax  ->Write();
    hBkg_phiDMin->Write();
    hBkg_phiDMax->Write();
    hBkg_etaDMin->Write();
    hBkg_etaDMax->Write();
    break;
  }
  // close file
  file->Close();
  delete file;
  // restore working area
  currentDir->cd();
  return;
}


void EventDiscriminator::read( const std::string& name ) {
  // save current working area
  TDirectory* currentDir = gDirectory;
  // open file
  TFile* file = new TFile( name.c_str() );
  // restore working area
  currentDir->cd();
  // read and copy histograms
  hSig_ptSum   = dynamic_cast<TH1F*>( file->Get( "hSig_ptSum"   )->Clone() );
  hSig_ptMin   = dynamic_cast<TH1F*>( file->Get( "hSig_ptMin"   )->Clone() );
  hSig_ptMax   = dynamic_cast<TH1F*>( file->Get( "hSig_ptMax"   )->Clone() );
  hSig_phiDMin = dynamic_cast<TH1F*>( file->Get( "hSig_phiDMin" )->Clone() );
  hSig_phiDMax = dynamic_cast<TH1F*>( file->Get( "hSig_phiDMax" )->Clone() );
  hSig_etaDMin = dynamic_cast<TH1F*>( file->Get( "hSig_etaDMin" )->Clone() );
  hSig_etaDMax = dynamic_cast<TH1F*>( file->Get( "hSig_etaDMax" )->Clone() );
  hBkg_ptSum   = dynamic_cast<TH1F*>( file->Get( "hBkg_ptSum"   )->Clone() );
  hBkg_ptMin   = dynamic_cast<TH1F*>( file->Get( "hBkg_ptMin"   )->Clone() );
  hBkg_ptMax   = dynamic_cast<TH1F*>( file->Get( "hBkg_ptMax"   )->Clone() );
  hBkg_etaDMin = dynamic_cast<TH1F*>( file->Get( "hBkg_etaDMin" )->Clone() );
  hBkg_phiDMax = dynamic_cast<TH1F*>( file->Get( "hBkg_phiDMax" )->Clone() );
  hBkg_phiDMin = dynamic_cast<TH1F*>( file->Get( "hBkg_phiDMin" )->Clone() );
  hBkg_etaDMax = dynamic_cast<TH1F*>( file->Get( "hBkg_etaDMax" )->Clone() );
  // close file
  file->Close();
  delete file;
  return;
}


// compute discriminating variable for an event
float EventDiscriminator::get( const Event& ev ) {
  // compute event variables
  compute( ev );
  // initialize probabilities
  float dSig = 1.0;
  float dBkg = 1.0;
  // update probabilities for each variable
  get( ptSum  , dSig, dBkg, hSig_ptSum  , hBkg_ptSum   );
  get( ptMin  , dSig, dBkg, hSig_ptMin  , hBkg_ptMin   );
  get( ptMax  , dSig, dBkg, hSig_ptMax  , hBkg_ptMax   );
  get( phiDMin, dSig, dBkg, hSig_phiDMin, hBkg_phiDMin );
  get( phiDMax, dSig, dBkg, hSig_phiDMax, hBkg_phiDMax );
  get( etaDMin, dSig, dBkg, hSig_etaDMin, hBkg_etaDMin );
  get( etaDMax, dSig, dBkg, hSig_etaDMax, hBkg_etaDMax );
  // return discriminating variable
  return dSig;
}


// normalize histogram to have s=1 with s=sum of bin content
void EventDiscriminator::normalizeHisto( TH1F* h ) {
  // initialize sum of contents
  float sum = 0;
  // loop over bins and sum contents
  int i;
  int n = h->GetNbinsX() + 2;
  for ( i = 0; i < n; ++i ) sum += h->GetBinContent( i );
  // loop over bins and divide content by the sum
  for ( i = 0; i < n; ++i ) {
    h->SetBinContent( i,
    h->GetBinContent( i ) / sum );
    h->SetBinError(   i,
    h->GetBinError(   i ) / sum );
  }
  return;
}


// compute event variables
void EventDiscriminator::compute( const Event& ev ) {

  // initialize variables
  ptSum   = 0.0;
  ptMin   = 999999.0;
  ptMax   = 0.0;
  phiDMin = 999999.0;
  etaDMin = 999999.0;
  phiDMax = 0.0;
  etaDMax = 0.0;
  // loop over particles
  const std::vector<Event::part_ptr> particles = ev.particles();
  int i;
  int j;
  int n = particles.size();
  // pt sum, pt min and pt max
  for ( i = 0; i < n; ++i ) {
    const Event::Particle& p = *particles[i];
    float pt = p.pt;
    ptSum += pt;
    if ( pt < ptMin ) ptMin = pt;
    if ( pt > ptMax ) ptMax = pt;
  }
  // phi and eta differences
  const float TWO_PI = 2 * M_PI;
  for ( i = 0; i < n; ++i ) {
    const Event::Particle& p = *particles[i];
    for ( j = i + 1; j < n; ++j ) {
      const Event::Particle& q = *particles[j];
      float etaDiff = fabs( p.eta - q.eta );
      float phiDiff = fabs( p.phi - q.phi );
      if ( phiDiff > M_PI ) phiDiff = TWO_PI - phiDiff;
      if ( etaDiff < etaDMin ) etaDMin = etaDiff;
      if ( etaDiff > etaDMax ) etaDMax = etaDiff;
      if ( phiDiff < phiDMin ) phiDMin = phiDiff;
      if ( phiDiff > phiDMax ) phiDMax = phiDiff;
    }
  }
  return;
}


// compute probabilities for all variables
void EventDiscriminator::get( float x,
                              float& dSig, float& dBkg,
                              TH1F*  hSig, TH1F*  hBkg ) {
  // get probabilities for signal and background
  float bkg = hBkg->GetBinContent( hBkg->FindBin( x ) );
  float sig = hSig->GetBinContent( hSig->FindBin( x ) );
  // skip zeros
  if ( sig < 1.0e-06 ) return;
  if ( bkg < 1.0e-06 ) return;
  // update probabilities
  dSig *= sig;
  dBkg *= bkg;
  // normalize probabilities
  float
  dSum = dBkg + dSig;
  dSig /= dSum;
  dBkg /= dSum;
  return;
}

