////////////////////////////////////////////////////////////////////
/// \file PlotFit.cc
///
/// \brief Functions to plot aspects of the fit, e.g. positions & energy
///
/// \author P G Jones <p.g.jones@qmul.ac.uk>
///
/// REVISION HISTORY:\n
///     2014-03-27 : P G Jones - First Revision.\n
///     2014-08-30 : I Coulter - Add additional functions showing how to access the reconstructed
///                              vertex for fitters of any name (rather than just the default)
///                              Also add more comments, descriptions and classifiers.
///     2018-04-17 : N Barros - Add the energy corrector into the example
///     2018-07-25 : Logan L  - Add the energy calibrator into the example
///
/// \details This file contains functions which show how to plot
///          aspects of the fit, such as energy, position, time or
///          direction as well as the results of a classifer.
///          These exist for the generic case in which the user wants
///          to specific the name of the fitter and for the special
///          case where the user just wants the last fitter applied
///          (for which functions are noted as Default).
///
/// Contains: PlotFitEnergy( fileName, fitName )         // These show how to plot the reconstructed energy, position,
///           PlotFitPosition( fileName, fitName )       // time or direction for a fitter with name "fitName"
///           PlotFitTime( fileName, fitName )
///           PlotFitDirection( fileName, fitName )
///
///           PlotClassification( fileName, className, classification )   // How to get the results of a classifier
///
///           PlotDefaultFitEnergy( fileName )           // These plot the reconstructed energy, position, time or
///           PlotDefaultFitPosition( fileName )         // direction for the defaultFitter within the rootfile,
///           PlotDefaultFitTime( fileName )             // That is, the last one to be applied
///           PlotDefaultFitDirectionPhi( fileName )
///           PlotDefaultFitDirectionTheta( fileName )
/// \details Functions to plot aspects of the fit.
///
////////////////////////////////////////////////////////////////////

#include <RAT/DU/DSReader.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/EV.hh>
#include <RAT/DS/PMT.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DB.hh>
#include <RAT/DU/Utility.hh>

#include <TH1D.h>
#include <TMath.h>
#include <TCanvas.h>

#include <string>
using namespace std;

/// Plot the fitted energy for a fitter with name "fitName"
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @param[in] fitName of the fitter whose results we want
void PlotFitEnergy( const string& fileName, const string& fitName )
{
  TH1D* hFitEnergy = new TH1D( "hFitEnergy", "Fit energy", 500, 0.0, 5.0 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);
  const RAT::DU::ReconCorrector &eCorr = RAT::DU::Utility::Get()->GetReconCorrector();
  const RAT::DU::ReconCalibrator &eCalib = RAT::DU::Utility::Get()->GetReconCalibrator();
  RAT::DU::DSReader dsReader( fileName );
  const RAT::DS::Run& run = dsReader.GetRun();
  const bool isData = !run.GetMCFlag();  // save whether data or simulation
  // Loop through entrys in rootfile
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

      // Look at each triggered event in the entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          // Get the EV
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );

          // Check for non-existant or poor fits
          if( !rEV.FitResultExists(fitName) ) continue;                             // No fit vertex exists
          if( !rEV.GetFitResult(fitName).GetVertex(0).ContainsEnergy() ) continue;  // Fit vertex doesn't contain an energy fit
          if( !rEV.GetFitResult(fitName).GetVertex(0).ValidEnergy() ) continue;     // Energy fit was marked as an invalid fit

          // Should be left with only the events with successful energy fits from the "fitName" fitter

          RAT::DS::FitResult fResult = rEV.GetFitResult(fitName);     // Get fit result
          RAT::DS::FitVertex fVertex = fResult.GetVertex(0);          // Get first fit vertex
          double fEnergy = fVertex.GetEnergy();                       // Get reconstructed energy
          //
          //
          // -- The correction below should only be applied for simulations produced using
          //    specific RAT releases:
          //
          // -- If using RAT versions 6.5.X the following correction should be used:
          //    fEnergy = eCorr.CorrectEnergyRSP(fEnergy,1);
          //
          // -- If using RAT versions 6.17.4 - 6.18.0, the following correction should be used:
          fEnergy = eCorr.CorrectEnergyRSP(fEnergy,2); // Correct the reconstructed energy
          //
          //
          // -- The calibration below was designed to work only with specific RAT versions and
          //    should be applied after the correction above (if the correction is needed).
          //
          // -- If using RAT versions 6.5.X the following calibration should be used:
          double x = fVertex.GetPosition().X();        // Get reconstructed position
          double y = fVertex.GetPosition().Y();
          double z = fVertex.GetPosition().Z();
          double rho = sqrt(x*x+y*y);
          //fEnergy = eCalib.CalibrateEnergyRSP( isData, fEnergy, rho, z, 0, 1 );    // Calibrate the reconstructed energy (after correction)
          //
          // -- If using RAT versions 6.17.4 - 6.18.7, the following calibration should be used:
          double xDir = fVertex.GetDirection().x();    // Direction info is needed only for version 2
          double yDir = fVertex.GetDirection().y();
          double zDir = fVertex.GetDirection().z();
          double UdotR = (xDir*x + yDir*y + zDir*z) / sqrt(rho*rho + z*z);
          fEnergy = eCalib.CalibrateEnergyRSP( isData, fEnergy, rho, z, UdotR, 2 );  // Calibrate the reconstructed energy (after correction)

          hFitEnergy->Fill( fEnergy );
        }
    }
  hFitEnergy->GetYaxis()->SetTitle( "Count per 10keV bin" );
  hFitEnergy->GetXaxis()->SetTitle( "Fitted energy [MeV]" );
  hFitEnergy->Draw();
}

/// Plot the fitted position for a fitter with name "fitName"
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @param[in] fitName of the fitter whose results we want
void PlotFitRadius( const string& fileName, const string& fitName )
{
  TH1D* hFitRadius = new TH1D( "hFitRadius", "Fit radius", 800, 0.0, 8000.0 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop through entrys in rootfile
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

      // Look at each triggered event in the entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          // Get the EV
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );

          // Check for non-existant or poor fits
          if( !rEV.FitResultExists(fitName) ) continue;                               // No fit vertex exists
          if( !rEV.GetFitResult(fitName).GetVertex(0).ContainsPosition() ) continue;  // Fit vertex doesn't contain a position fit
          if( !rEV.GetFitResult(fitName).GetVertex(0).ValidPosition() ) continue;     // Position fit was marked as an invalid fit

          // Should be left with only the events with successful position fits from the "fitName" fitter

          RAT::DS::FitResult fResult = rEV.GetFitResult(fitName);     // Get fit result
          RAT::DS::FitVertex fVertex = fResult.GetVertex(0);          // Get first fit vertex
          TVector3 fPosition = fVertex.GetPosition();                 // Get reconstructed position
          hFitRadius->Fill( fPosition.Mag() );
        }
    }
  hFitRadius->GetYaxis()->SetTitle( "Count per 10mm bin" );
  hFitRadius->GetXaxis()->SetTitle( "Fitted radius [mm]" );
  hFitRadius->Draw();
}

/// Plot the fitted time for a fitter with name "fitName"
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @param[in] fitName of the fitter whose results we want
void PlotFitTime( const string& fileName, const string& fitName )
{
  TH1D* hFitTime = new TH1D( "hFitTime", "Fit time", 100, 0, 400 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop through entrys in rootfile
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

      // Look at each triggered event in the entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          // Get the EV
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );

          // Check for non-existant or poor fits
          if( !rEV.FitResultExists(fitName) ) continue;                           // No fit vertex exists
          if( !rEV.GetFitResult(fitName).GetVertex(0).ContainsTime() ) continue;  // Fit vertex doesn't contain a time fit
          if( !rEV.GetFitResult(fitName).GetVertex(0).ValidTime() ) continue;     // Time fit was marked as an invalid fit

          // Should be left with only the events with successful time fits from the "fitName" fitter

          RAT::DS::FitResult fResult = rEV.GetFitResult(fitName);     // Get fit result
          RAT::DS::FitVertex fVertex = fResult.GetVertex(0);          // Get first fit vertex
          double fTime = fVertex.GetTime();                           // Get reconstructed time
          hFitTime->Fill( fTime );
        }
    }
  hFitTime->GetYaxis()->SetTitle( "Count per 4ns bin" );
  hFitTime->GetXaxis()->SetTitle( "Fitted time [ns]" );
  hFitTime->Draw();
}

/// Plot the fitted direction for a fitter with name "fitName"
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @param[in] fitName of the fitter whose results we want
void PlotFitDirection( const string& fileName, const string& fitName )
{
  TH1D* hFitDirectionPhi = new TH1D( "hFitDirectionPhi", "Fit direction (phi)", 180, -180, 180 );
  TH1D* hFitDirectionTheta = new TH1D( "hFitDirectionTheta", "Fit direction (theta)", 90, 0, 180 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop through entrys in rootfile
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

      // Look at each triggered event in the entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          // Get the EV
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );

          // Check for non-existant or poor fits
          if( !rEV.FitResultExists(fitName) ) continue;                                // No fit vertex exists
          if( !rEV.GetFitResult(fitName).GetVertex(0).ContainsDirection() ) continue;  // Fit vertex doesn't contain a direction fit
          if( !rEV.GetFitResult(fitName).GetVertex(0).ValidDirection() ) continue;     // Direction fit was marked as an invalid fit

          // Should be left with only the events with successful direction fits from the "fitName" fitter

          RAT::DS::FitResult fResult = rEV.GetFitResult(fitName);               // Get fit result
          RAT::DS::FitVertex fVertex = fResult.GetVertex(0);                    // Get first fit vertex
          TVector3 fDirection = fVertex.GetDirection();                         // Get reconstructed direction
          hFitDirectionPhi->Fill( fDirection.Phi()*TMath::RadToDeg() );
          hFitDirectionTheta->Fill( fDirection.Theta()*TMath::RadToDeg() );
        }
    }
  TCanvas* can = new TCanvas("cDir","Direction Plots", 600, 600);
  can->Divide(1,2);
  can->cd(1);
  hFitDirectionPhi->GetYaxis()->SetTitle( "Count per 2 degree bin" );
  hFitDirectionPhi->GetXaxis()->SetTitle( "Fitted phi direction [#circ]" );
  hFitDirectionPhi->Draw();
  can->cd(2);
  hFitDirectionTheta->GetYaxis()->SetTitle( "Count per 2 degree bin" );
  hFitDirectionTheta->GetXaxis()->SetTitle( "Fitted theta direction [#circ]" );
  hFitDirectionTheta->Draw();
}

/// Plot the results of "classification" from the classifier "className"
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @param[in] className of the classifier whose results we want
/// @param[in] classication is the type result we want from the classifier
void PlotClassification( const string& fileName, const string& className, const string& classification )
{
  // Range of output values with depend on the classifier
  // In this example, assume they range between 0 and 1 (as might be found in beta14)
  TH1D* hClass = new TH1D( "hClass", "Classification", 100, 0.0, 1.0 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop through entrys in rootfile
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

      // Look at each triggered event in the entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          // Get the EV
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );

          // Check for non-existant or poor classifications
          if( !rEV.ClassifierResultExists(className) ) continue;          // No classifier result exists
          if( !rEV.GetClassifierResult(className).GetValid() ) continue;  // Classifier result is marked as invalid

          // Should be left with only the events with successful classifications from the "className" classifier

          RAT::DS::ClassifierResult cResult = rEV.GetClassifierResult(className);     // Get classifier result
          double cValue = cResult.GetClassification(classification);                       // Get the value of "classification" from the classifier
          hClass->Fill(cValue);
        }
    }
  hClass->GetYaxis()->SetTitle( "Counts" );
  hClass->GetXaxis()->SetTitle( "Classifier value" );
  hClass->Draw();
}


/// Plot the fitted energy for the defaultFitter
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotDefaultFitEnergy( const string& fileName )
{
  TH1D* hFitEnergy = new TH1D( "hFitEnergy", "Fit energy", 500, 0.0, 5.0 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop over entries
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      // For each entry, loop over triggered events
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
	  // Check that fit exists, contains an energy fit and that energy fit was successful
          if( !rEV.DefaultFitVertexExists() || !rEV.GetDefaultFitVertex().ContainsEnergy() || !rEV.GetDefaultFitVertex().ValidEnergy() )
            continue; // Didn't fit successfully
          hFitEnergy->Fill( rEV.GetDefaultFitVertex().GetEnergy() );
        }
    }
  hFitEnergy->GetYaxis()->SetTitle( "Count per 10keV bin" );
  hFitEnergy->GetXaxis()->SetTitle( "Fitted energy [MeV]" );
  hFitEnergy->Draw();
  return hFitEnergy;
}

/// Plot the fitted radial position for the defaultFitter
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotDefaultFitRadius( const string& fileName )
{
  TH1D* hFitRadius = new TH1D( "hFitRadius", "Fit radius", 800, 0.0, 8000.0 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop over entries
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      // Loop over triggered events within each entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
          // Check that fit exists, contains a position fit and that position fit was successful
          if( !rEV.DefaultFitVertexExists() || !rEV.GetDefaultFitVertex().ContainsPosition() || !rEV.GetDefaultFitVertex().ValidPosition() )
            continue; // Didn't fit successfully
          hFitRadius->Fill( rEV.GetDefaultFitVertex().GetPosition().Mag() );
        }
    }
  hFitRadius->GetYaxis()->SetTitle( "Count per 10mm bin" );
  hFitRadius->GetXaxis()->SetTitle( "Fitted radius [mm]" );
  hFitRadius->Draw();
  return hFitRadius;
}

/// Plot the fitted time for the defaultFitter
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotDefaultFitTime( const string& fileName )
{
  TH1D* hFitTime = new TH1D( "hFitTime", "Fit time", 100, 0, 400 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop over entries
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      // Loop over triggered events for each entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
          // Check that fit exists, contains a time fit and that time fit was successful
          if( !rEV.DefaultFitVertexExists() || !rEV.GetDefaultFitVertex().ContainsTime() || !rEV.GetDefaultFitVertex().ValidTime() )
            continue; // Didn't fit successfully
          hFitTime->Fill( rEV.GetDefaultFitVertex().GetTime() );
        }
    }
  hFitTime->GetYaxis()->SetTitle( "Count per 4ns bin" );
  hFitTime->GetXaxis()->SetTitle( "Fitted time [ns]" );
  hFitTime->Draw();
  return hFitTime;
}

/// Plot the fitted direction phi
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotDefaultFitDirectionPhi( const string& fileName )
{
  TH1D* hFitDirectionPhi = new TH1D( "hFitDirectionPhi", "Fit direction (phi)", 180, -180, 180 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop over entries
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      // Loop over triggered events for each entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
          // Check that fit exists, contains a direction fit and that direction fit was successful
          if( !rEV.DefaultFitVertexExists() || !rEV.GetDefaultFitVertex().ContainsDirection() || !rEV.GetDefaultFitVertex().ValidDirection() )
            continue; // Didn't fit successfully
          hFitDirectionPhi->Fill( rEV.GetDefaultFitVertex().GetDirection().Phi() * TMath::RadToDeg() );
        }
    }
  hFitDirectionPhi->GetYaxis()->SetTitle( "Count per 2 degree bin" );
  hFitDirectionPhi->GetXaxis()->SetTitle( "Fitted phi direction [#circ]" );
  hFitDirectionPhi->Draw();
  return hFitDirectionPhi;
}

/// Plot the fitted direction theta
///
/// @param[in] fileName of the RAT::DS root file to analyse
/// @return the histogram plot
TH1D* PlotDefaultFitDirectionTheta( const string& fileName )
{
  TH1D* hFitDirectionTheta = new TH1D( "hFitDirectionTheta", "Fit direction (theta)", 90, 0, 180 );
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader( fileName );
  // Loop over entries
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );
      // Loop over triggered events within entry
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
          // Check that fit exists, contains a direction fit and that direction fit was successful
          if( !rEV.DefaultFitVertexExists() || !rEV.GetDefaultFitVertex().ContainsDirection() || !rEV.GetDefaultFitVertex().ValidDirection() )
            continue; // Didn't fit successfully
          hFitDirectionTheta->Fill( rEV.GetDefaultFitVertex().GetDirection().Theta() * TMath::RadToDeg() );
        }
    }
  hFitDirectionTheta->GetYaxis()->SetTitle( "Count per 2 degree bin" );
  hFitDirectionTheta->GetXaxis()->SetTitle( "Fitted theta direction [#circ]" );
  hFitDirectionTheta->Draw();
  return hFitDirectionTheta;
}