#include <TH1D.h>
#include <TH2D.h>
#include <TFile.h>
#include <TMath.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TLegend.h>

// -- RAT headers
#include <RAT/DU/DSReader.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/MC.hh>
#include <RAT/DS/MCPE.hh>
#include <RAT/DS/MCPMT.hh>
#include <RAT/DU/LightPathCalculator.hh>

#include <iostream>
#include <string>

using std::cout;
using std::endl;
using std::string;

/** These functions are using the MC branches and are looking for information
that is totally stored in the ROOT file.
Therefore, there is no need for the database connection, in which case we can
spare the server a few more cycles.
The function SetupFullOfflineMode shows a way of how to do that.
**/

///
/// WARNING: Setting up full ofline mode is dangerous. You could be missing important
///          information from the database leading to wrong results.
///          This should only be used when people really know what they are doing and
///          are *ABSOLUTELY SURE* that there is no relevant information in the database.
///

void SetupFullOfflineMode() {
  RAT::DB *db = RAT::DB::Get();
  // disable remote database connection
  db->SetAirplaneModeStatus(true);
  // Disable default plane locks
  db->SetDefaultPlaneLockStatus(false);
}

/// Plot the number of MC photoelectrons per event (or NumPE)
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotMCPhotoelectronNhit( const string& fileName )
{
  // Before anything else is done. Disable the database
  SetupFullOfflineMode();
  TH1D* hNumPE = new TH1D( "hNumPE", "Number of photoelectrons per event", 2000, 0.0, 2000.0 );

  RAT::DU::DSReader dsReader( fileName );
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      if (!rDS.MCExists()) {
        // There is no MC. Stop the loop
        break;
      }
      hNumPE->Fill( rDS.GetMC().GetMCPECount() );
    }
  hNumPE->GetYaxis()->SetTitle( "Count per 1 pe bin" );
  hNumPE->GetXaxis()->SetTitle( "Number of photoelectrons per event" );
  hNumPE->Draw();
  return hNumPE;
}

/// Plot the number of MC hits per event
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotMCHitsNhit( const string& fileName )
{
  // Before anything else is done. Disable the database
  SetupFullOfflineMode();
  TH1D* hNumHits = new TH1D( "hNumHits", "Number of MC hits per event", 2000, 0.0, 2000.0 );

  RAT::DU::DSReader dsReader( fileName );
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      if (!rDS.MCExists()) {
        // There is no MC. Stop the loop
        break;
      }
      for( size_t iMCEvent = 0; iMCEvent < rDS.GetMCEVCount(); iMCEvent++ ) // Loop over triggered events
        hNumHits->Fill( rDS.GetMCEV( iMCEvent ).GetMCHits().GetAllCount() ); // All can be changed to Normal, OWL etc...
    }
  hNumHits->GetYaxis()->SetTitle( "Count per 1 hit bin" );
  hNumHits->GetXaxis()->SetTitle( "Number of MC hits per event" );
  hNumHits->Draw();
  return hNumHits;
}

/// Plot the number of Uncalibrated hits per event
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotUncalibratedNhit( const string& fileName )
{
  // Before anything else is done. Disable the database
  SetupFullOfflineMode();

  TH1D* hUncalHits = new TH1D( "hUncalHits", "Number of Uncalibrated hits per event", 2000, 0.0, 2000.0 );

  RAT::DU::DSReader dsReader( fileName );
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        hUncalHits->Fill( rDS.GetEV( iEV ).GetUncalPMTs().GetAllCount() );
    }
  hUncalHits->GetYaxis()->SetTitle( "Count per 1 hit bin" );
  hUncalHits->GetXaxis()->SetTitle( "Number of Uncalibrated hits per event" );
  hUncalHits->Draw();
  return hUncalHits;
}

/// Plot the number of Calibrated hits per event
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotCalibratedNhit( const string& fileName )
{
    // Before anything else is done. Disable the database
  SetupFullOfflineMode();

  TH1D* hCalHits = new TH1D( "hCalHits", "Number of Calibrated hits per event", 2000, 0.0, 2000.0 );

  RAT::DU::DSReader dsReader( fileName );
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        hCalHits->Fill( rDS.GetEV( iEV ).GetCalPMTs().GetAllCount() );
    }
  hCalHits->GetYaxis()->SetTitle( "Count per 1 hit bin" );
  hCalHits->GetXaxis()->SetTitle( "Number of calibrated hits per event" );
  hCalHits->Draw();
  return hCalHits;
}

/// Plot all NHit numbers
///
/// @param[in] fileName of the RAT::DS root file to analyse
void PlotNhit( const string& fileName )
{
  // Before anything else is done. Disable the database
  SetupFullOfflineMode();

  TCanvas* c1 = new TCanvas();
  cout << "PlotNhit : Drawing MCPhotoelectron" << endl;
  TH1D* numPE = NULL;
  numPE = PlotMCPhotoelectronNhit( fileName );
  cout << "PlotNhit : Drawing MCHits" << endl;
  TH1D* mcHits = PlotMCHitsNhit( fileName );
  cout << "PlotNhit : Drawing Uncal NHit" << endl;
  TH1D* uncalibrated = PlotUncalibratedNhit( fileName );
  cout << "PlotNhit : Drawing Cal NHit" << endl;
  TH1D* calibrated = PlotCalibratedNhit( fileName );
  numPE->Draw();
  mcHits->SetLineColor( kGreen + 2 );
  mcHits->Draw("SAME");
  uncalibrated->SetLineColor( kBlue );
  uncalibrated->Draw("SAME");
  calibrated->SetLineColor( kRed );
  calibrated->Draw("SAME");
  TLegend* t1 = new TLegend( 0.7, 0.7, 0.9, 0.9 );
  t1->AddEntry( numPE, "Photoelectrons", "l" );
  t1->AddEntry( mcHits, "MC Hits", "l" );
  t1->AddEntry( uncalibrated, "Uncalibrated", "l" );
  t1->AddEntry( calibrated, "Calibrated", "l" );
  t1->Draw();
  c1->Update();
}