///////////////////////////////////////////////////////////////////////////////
/// \file protondEdx.cc
///
/// \brief Extracts dE/dx from tracking
///
/// \author P G Jones <p.g.jones@qmul.ac.uk>
///
/// REVISION HISTORY:\n
///     2014-07-21 : P G Jones - Moved to rat-tools validation.\n
///     2016-10-19 : M Stringer - Changes to how MCTracks are accessed (PR #1508).
//      2018-02-10 : R Lane - change to function arguments necessary for ROOT 6 compatibility
///
/// \details The dEdx at various energies is extracted
///
///////////////////////////////////////////////////////////////////////////////
#include <RAT/DU/DSReader.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/MC.hh>
#include <RAT/DS/MCTrack.hh>
#include <RAT/DS/MCTrackStep.hh>

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

void protondEdx(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");

  const double energyBins[] = { 0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95,
                                1.50, 2.50, 3.50, 4.50, 5.50, 6.50, 7.50, 8.50, 9.50,
                                15.0, 25.0, 35.0, 45.0, 55.0, 65.0, 75.0, 85.0, 95.0, 200.0 };
  TH1D* dedx = new TH1D( "dedx", "dedx", 28, energyBins );
  TH1D* count = new TH1D( "count", "count", 28, energyBins );
  dedx->SetXTitle( "Energy in MeV." );
  dedx->SetYTitle( "dE/dx in MeV/mm." );

  for( size_t iEntry = 0; iEntry < dsReader->GetEntryCount(); iEntry++ )
    {
      const RAT::DS::Entry& rDS = dsReader->GetEntry( iEntry );
      const RAT::DS::MC& rMC = rDS.GetMC();
      std::vector<unsigned int> mcTrackIDs = rMC.GetMCTrackIDs();
      for( size_t iTrack = 0; iTrack < mcTrackIDs.size(); iTrack++ )
        {
          const RAT::DS::MCTrack& rMCTrack = rMC.GetMCTrack( mcTrackIDs.at(iTrack) );
          if( rMCTrack.GetParentID() != 0 ) // Primary track has parent ID 0
            continue;
          double previousEnergy = rMCTrack.GetMCTrackStep( 0 ).GetKineticEnergy();
          for( size_t iStep = 1; iStep < rMCTrack.GetMCTrackStepCount(); iStep++ )
            {
              const RAT::DS::MCTrackStep& rMCTrackStep = rMCTrack.GetMCTrackStep( iStep );
              const double dE = previousEnergy - rMCTrackStep.GetKineticEnergy();
              if( dE > 0.0 && rMCTrackStep.GetLength() > 0.0 )
                {
                  dedx->Fill( previousEnergy, dE / rMCTrackStep.GetLength() );
                  count->Fill( previousEnergy );
                }
              previousEnergy = rMCTrackStep.GetKineticEnergy();
            }
        }
    }
  dedx->Divide( count );

  outtfile->cd();
  dedx->Write();
  outtfile->Close();
  delete outtfile;
  delete dsReader;
}