////////////////////////////////////////////////////////////////////
/// \file waterfitter.cc
///
/// \brief Extracts the reconstructed position, energy and direction
///
/// \author I Coulter <icoulter@hep.upenn.edu>
///
/// REVISION HISTORY:\n
///     2014-05-29 : P G Jones - Added header information.\n
//      2018-02-10 : R Lane - change to function arguments necessary for ROOT 6 compatibility
///
/// \details The reconstructed position energy and direction are
/// plotted against each other and MC.
///
////////////////////////////////////////////////////////////////////
#include <RAT/DU/DSReader.hh>
#include <RAT/DS/Entry.hh>

#include <TH1D.h>
#include <TFile.h>
#include <TProfile.h>
#include <TVector3.h>

void waterfitter(std::string event_file, std::string outfile)
{
  RAT::DU::DSReader* dsReader = new RAT::DU::DSReader(event_file);
  TFile *outtfile = new TFile(outfile.c_str(),"RECREATE");

  TH1D *hCenterRes = new TH1D("hCenterRes", "6 MeV electrons at the center;Reconstructed Radius (mm);Events per bin", 10, 0, 1000);
  TProfile *hReconVsRadius = new TProfile("hReconVsRadius", "6 MeV electrons (error bars are RMS);(Radius/6005 mm)**3;Reconstructed Pos - True Pos (mm)", 5, 0, 1, 0, 750.0, "s");
  TProfile *hReconVsEnergy = new TProfile("hReconVsEnergy", "Electrons at center (error bars are RMS);Kinetic energy (MeV);Reconstructed Pos (mm)", 5, 5.9, 10.1, 0, 750.0, "s");
  TH1D *hCenterDirection = new TH1D("hCenterDir", "6 MeV electrons at the center;Cosine of angle between reconstructed direction and MC diretion;Events per bin", 20, -1, 1);
  TProfile *hDirectionVsRadius = new TProfile("hDirectionVsRadius", "6 MeV electrons (error bars are RMS);(Radius/6005 mm)**3;Cosine of angle between reconstructed direction and MC diretion", 5, 0, 1, -1, 1, "s");
  TProfile *hDirectionVsEnergy = new TProfile("hDirectionVsEnergy", "Electrons at center (error bars are RMS);Kinetic Energy (MeV); Cosine of angle between reconstructed direction and MC diretion", 5, 5.9, 10.1, -1, 1, "s");
  TProfile *hEnergyVsRadius = new TProfile("hEnergyReconVsRadius", "6MeV Electrons at center (error bars are RMS);(Radius/6005 mm)**3;Reconstructed energy - True energy (MeV)", 5, 0, 1, -1.0, 1.0, "s");
  TProfile *hEnergyVsEnergy = new TProfile("hEnergyReconVsEnergy", "Electrons at center (error bars are RMS);Kinetic energy (MeV);Reconstructed energy - True energy (MeV)", 5, 5.9, 10.1, -1.0, 1.0, "s");

  for( size_t iEntry = 0; iEntry < dsReader->GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader->GetEntry( iEntry );
      const TVector3 mcPos = rDS.GetMC().GetMCParticle(0).GetPosition();
      const TVector3 mcDir = rDS.GetMC().GetMCParticle(0).GetMomentum().Unit();
      const double mcEnergy = rDS.GetMC().GetMCParticle(0).GetKineticEnergy();
      for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ )
        {
          const RAT::DS::EV& rEV = rDS.GetEV( iEV );
          if( !rEV.FitResultExists("waterFitter") || rEV.GetFitResult("waterFitter").GetVertexCount() == 0
              || !rEV.GetFitResult("waterFitter").GetValid() )
            continue;

          const TVector3 fitPos = rEV.GetFitResult("waterFitter").GetVertex(0).GetPosition();
          const TVector3 fitDir = rEV.GetFitResult("waterFitter").GetVertex(0).GetDirection();
          const double ctheta = fitDir.Dot(mcDir);

          if( mcEnergy < 6.05 && mcPos.Mag() < 0.5 )
            hCenterRes->Fill( fitPos.Mag() );
          if( mcEnergy > 6.05 && mcPos.Mag() < 0.5 )
            hReconVsEnergy->Fill( mcEnergy, fitPos.Mag() );
          if( mcEnergy < 6.05 && mcPos.Mag() > 0.5 )
            hReconVsRadius->Fill( pow( mcPos.Mag() / 6005.0, 3.0 ), (fitPos - mcPos).Mag() );

          if( mcEnergy < 6.05 && mcPos.Mag() < 0.5 )
            hCenterDirection->Fill(ctheta);
          if( mcEnergy > 6.05 && mcPos.Mag() < 0.5 )
            hDirectionVsEnergy->Fill( mcEnergy, ctheta);
          if( mcEnergy < 6.05 && mcPos.Mag() > 0.5 )
            hDirectionVsRadius->Fill( pow( mcPos.Mag() / 6005.0, 3.0 ), ctheta );

          // Get the energy fit
          double fitEnergy = rEV.GetFitResult("waterFitter").GetVertex(0).GetEnergy();
          if( mcPos.Mag() < 0.5 )
            hEnergyVsEnergy->Fill(mcEnergy, fitEnergy-mcEnergy);
          if( mcEnergy < 6.05 && mcPos.Mag() > 0.5 )
            hEnergyVsRadius->Fill( pow(mcPos.Mag()/6005.0, 3.0), fitEnergy-mcEnergy);
        }
    }

  hReconVsRadius->SetMaximum(600);
  hReconVsRadius->SetMinimum(0);
  hReconVsEnergy->SetMaximum(600);
  hReconVsEnergy->SetMinimum(0);
  hEnergyVsEnergy->SetMaximum(1);
  hEnergyVsEnergy->SetMinimum(-1);
  hEnergyVsRadius->SetMaximum(1);
  hEnergyVsRadius->SetMinimum(-1);
  hDirectionVsEnergy->SetMaximum(1);
  hDirectionVsEnergy->SetMinimum(-1);
  hDirectionVsRadius->SetMaximum(1);
  hDirectionVsRadius->SetMinimum(-1);

  outtfile->cd();
  hCenterRes->Write();
  hReconVsRadius->Write();
  hReconVsEnergy->Write();
  hCenterDirection->Write();
  hDirectionVsRadius->Write();
  hDirectionVsEnergy->Write();
  hEnergyVsEnergy->Write();
  hEnergyVsRadius->Write();
  outtfile->Close();
  delete outtfile;
  delete dsReader;
}