ProtoDUNETriggeredBeam_module.cc

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
// Class:       ProtoDUNETriggeredBeam
// Module Type: producer
// File:        ProtoDUNETriggeredBeam_module.cc
//
// Modified from the original ProtoDUNEBeam module to more accurately
// reflect what happens in reality. As such, it does not make use of
// the GoodParticle tree and uses the TRIG planes to work out which
// events are "Triggered" events
//
// Leigh Whitehead & Jake Calcutt
//
////////////////////////////////////////////////////////////////////////

#include "art/Framework/Core/EDProducer.h"
#include "art/Framework/Core/ModuleMacros.h"
#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/Handle.h"
#include "art/Framework/Principal/Run.h"
#include "art/Framework/Principal/SubRun.h"
#include "canvas/Utilities/InputTag.h"
#include "fhiclcpp/ParameterSetRegistry.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"
#include "cetlib_except/exception.h"
#include "cetlib/search_path.h"
#include "cetlib/filesystem.h"
#include "nusimdata/SimulationBase/MCTruth.h"
#include "nusimdata/SimulationBase/MCParticle.h"
#include "larcore/Geometry/Geometry.h"
#include "larcoreobj/SummaryData/RunData.h"
#include <memory>
#include <string>
#include <map>
#include <utility>
#include <vector>
#include "dunecore/DuneObj/ProtoDUNEBeamEvent.h"
// art extensions
#include "nurandom/RandomUtils/NuRandomService.h"
#include "art_root_io/TFileService.h"
#include <TSystem.h>
#include <TFile.h>
#include <TTree.h>
#include <TGraph.h>
#include <TVector3.h>
#include "THnSparse.h"
#include "TH2D.h"
#include "TRandom3.h"
#include <TLorentzVector.h>
#include <TDatabasePDG.h>
#include <TParticlePDG.h>
#include "CLHEP/Random/RandFlat.h"
#include "ifdh.h"
#include <sys/stat.h>

namespace evgen{

    class ProtoDUNETriggeredBeam;
 
    class ProtoDUNETriggeredBeam : public art::EDProducer {
    public:
        explicit ProtoDUNETriggeredBeam(fhicl::ParameterSet const & p);
        // The destructor generated by the compiler is fine for classes
        // without bare pointers or other resource use.
        ~ProtoDUNETriggeredBeam();
        
        // Plugins should not be copied or assigned.
        ProtoDUNETriggeredBeam(ProtoDUNETriggeredBeam const &) = delete;
        ProtoDUNETriggeredBeam(ProtoDUNETriggeredBeam &&) = delete;
        ProtoDUNETriggeredBeam & operator = (ProtoDUNETriggeredBeam const &) = delete;
        ProtoDUNETriggeredBeam & operator = (ProtoDUNETriggeredBeam &&) = delete;
        
        // Required functions.
        void produce(art::Event & e) override;
        void beginJob() override;
        void beginRun(art::Run& run) override;
        void endJob() override;

        // Simple struct to store the information for each particle at the front face
        struct BeamParticle {
          BeamParticle(){
            fTrackID=-999;
            fPDG=-999;
            fParentID=-999;
            fPosX = -999;
            fPosY = -999;
            fPosZ = -999;
            fPosT = -999;
            fMomX = -999;
            fMomY = -999;
            fMomZ = -999; 
          };
          BeamParticle(int trackid, int pdg, int parentid, float posX, float posY, float posZ, float posT,
                       float momX, float momY, float momZ){
            fTrackID = trackid;
            fPDG     = pdg;
            fParentID= parentid;
            fPosX    = posX;
            fPosY    = posY;
            fPosZ    = posZ;  
            fPosT    = posT;
            fMomX    = momX;
            fMomY    = momY;
            fMomZ    = momZ;
          };
          void Print(){
            std::cout << "Particle " << fPDG << ": (" << fPosX << "," << fPosY << "," << fPosZ << "," << fPosT << ") "
                                     << ": (" << fMomX << "," << fMomY << "," << fMomZ << ") " << std::endl;
          };
          int fTrackID, fPDG, fParentID;
          float fPosX, fPosY, fPosZ, fPosT;
          float fMomX, fMomY, fMomZ;
        };

        // Struct to contain the particles reaching the cryostat wall for each event
        // in the beam simulation files
        struct BeamEvent {
          BeamEvent(){
            fEventID = -999;
            fTriggerID = -999;
            fHasInteracted = false;
          }
          BeamEvent(int eventid){
            fEventID = eventid;
            fTriggerID = -999;
            fHasInteracted = false;
          };

          void AddParticle(BeamParticle particle){
            fParticlesFront.insert(std::make_pair(particle.fTrackID,particle));
          };

          int fEventID;

          // Map of particles to the track ID
          std::map<int,BeamParticle> fParticlesFront;

          int fTriggerID;

          // We need information for each point in the beamline
          std::map<std::string,BeamParticle> fTriggeredParticleInfo;

          // Some events can interact before between TRIG2 and NP04front
          bool fHasInteracted;
          std::vector<int> fSecondaryTrackIDs;
        };

        // Convenience struct to encapsulate all particles that would
        // deposit energy within one readout window of the TPC
        struct OverlaidTriggerEvent {

          OverlaidTriggerEvent(int trigID){
            fTriggerEventID = trigID;
          };

          void AddOverlay(int overlayID){
            fOverlayEventIDs.push_back(overlayID);
          };

          std::vector<int> fOverlayEventIDs;
          int fTriggerEventID;

        };
        
    private:

        // Calculate how many overlay events we need.
        void CalculateNOverlays();
        
        // Fill the above maps and vector.
        void FillParticleMaps(TTree *frontFaceTree);
       
        // Find trigger events
        std::vector<int> FindTriggeredEvents(TTree *trig1Tree,TTree *trig2Tree);
 
        // Add the triggered particle information for a given instrument
        void FillInstrumentInformation(std::vector<int> &eventIDs,TTree *instrumentTree);

        // Group the events to overlay a number of background events on each trigger event
        OverlaidTriggerEvent GenerateOverlaidEvent(const int &trigEventID);

        // Generate a true event based on a single entry from the input tree.
        void GenerateTrueEvent(simb::MCTruth &mcTruth, const OverlaidTriggerEvent &overlayEvent, beam::ProtoDUNEBeamEvent & beamEvent);

        // Convert our BeamParticle struct into a MCParticle object
        simb::MCParticle BeamParticleToMCParticle(const BeamParticle &beamParticle, const int outputTrackID, const float triggerParticleTime, const int primaryStatus, const std::string process, const int motherID = -1);

        // Use DDMC to create a primary beam particle
        simb::MCParticle DataDrivenMCParticle(
            const BeamParticle &beamParticle, const int outputTrackID,
            const float triggerParticleTime,
            beam::ProtoDUNEBeamEvent & beamEvent, const int primaryStatus, const std::string process);

        void SetDataDrivenPosMom(TLorentzVector & position,
                                 TLorentzVector & momentum, double sampledHUp,
                                 double sampledVUp, double sampledHDown,
                                 double sampledVDown, double sampledMomentum,
                                 int beamPDG);

        double UnsmearMomentum2D(double momentum, int pdg);

        void SetDataDrivenBeamEvent(
            beam::ProtoDUNEBeamEvent & beamEvent,
            double sampledHUp, double sampledVUp, double sampledHDown,
            double sampledVDown, double sampledMomentum);

        void ConvertSamplingPoint(double input_point[5],
                                  std::vector<double> minima,
                                  std::vector<double> maxima,
                                  double output_point[5]);

        // Handle root files from beam instrumentation group
        void OpenInputFile(std::string & filename);
        
        // Convert to the detector coordinate frame
        void ConvertCoordinates(float &x, float &y, float &z);
        
        // Convert the momentum to GeV and rotate as required.
        void ConvertMomentum(float &px, float &py, float &pz);

        // We need to rotate the beam monitor coordinates into the detector frame
        TLorentzVector ConvertBeamMonitorCoordinates(float x, float y, float z, float t, float offset);

        // Methods for making beam instrument tracks
        TVector3 ConvertProfCoordinates(double x, double y, double z, double zOffset);
        TVector3 ProjectToTPC(TVector3 firstPoint, TVector3 secondPoint);
        double GetPosition( short fiber );
        void MakeTracks( beam::ProtoDUNEBeamEvent & beamEvent );
        void MomentumSpectrometer( beam::ProtoDUNEBeamEvent & beamEvent );
        double MomentumCosTheta( double, double, double );
        std::string FindFile(const std::string filename);

        TVector3 ConvertBeamMonitorMomentumVec(float px, float py, float pz);
        // Setup the beam monitor basis vectors in detector coordinates
        void BeamMonitorBasisVectors();
        // Apply the rotation
        void RotateMonitorVector(TVector3 &vec);
        // Fill all of the ProtoDUNEBeamEvent information from the beam simulation trigger event
        void SetBeamEvent(beam::ProtoDUNEBeamEvent & beamevt, const BeamEvent &triggerEvent);
        beam::FBM MakeFiberMonitor( float pos );

        // Background particles need to be fired from an upstream position
        void SetBackgroundPosition(BeamParticle &particle);

        // Variables       
        CLHEP::RandFlat fFlatRnd;

        // Beam input file name and tree names (in beam order)
        std::string fFileName;
        std::string fBaseFileName;
        std::string fTOF1TreeName;
        std::string fBPROF1TreeName;
        std::string fBPROF2TreeName;
        std::string fBPROF3TreeName;
        std::string fTRIG1TreeName;
        std::string fBPROFEXTTreeName;
        std::string fBPROF4TreeName;
        std::string fTRIG2TreeName;
        std::string fNP04frontTreeName;

        // Variables for the beam simulation event building        
        std::vector<OverlaidTriggerEvent> fAllOverlaidTriggerEvents;
        std::map<int,BeamEvent> fAllBeamEvents;
        std::vector<int> fFinalTriggerEventIDs;
       
        // This event refers to the position in the list of beam simulation events 
        int fEventNumber;
        
        // Let the user define the event to start at
        int fStartEvent;

        // An important feature is to be able to use a data-driven triggered particle
        bool fUseDataDriven;
        
        // Define the coordinate transform from the beam frame to the detector frame
        float fBeamX;
        float fBeamY;
        float fBeamZ;
        float fBeamThetaShift;
        float fBeamPhiShift;
        // Rotate the beam monitor coordinate system (those after the last bending magnet)
        float fRotateMonitorXZ;
        float fRotateMonitorYZ;
        // The three beam monitor basis vectors in the detector coordinate system
        TVector3 fBMBasisX; 
        TVector3 fBMBasisY; 
        TVector3 fBMBasisZ; 
        // The z positions of the important elements along the beam direction
        float fBPROFEXTPos;
        float fBPROF4Pos;
        float fNP04frontPos;
        float fTRIG2Pos;
       
        // Parameters from the .fcl file to deal with overlaying events
        float fIntensity; // Number of interactions on the secondary target per SPS spill
        float fReadoutWindow; // Readout window (needs to match the values used in the simulation) in milliseconds
        float fBeamSpillLength; // The SPS spill length in seconds

        float fLB;
        float fL1;
        float fL2;
        float fL3;
        float fBeamBend;

        float fLMag;
        std::string fNominalP;
        float fB;

        int fMaxSamples;

        TRandom3 fRNG;
        bool fVerbose, fIncludeAnti;
        std::vector<double> fMinima = {0.8, 0., 0., 0., 0.};
        std::vector<double> fMaxima = {1.2, 192., 192., 192., 192.};
        std::map<int, std::string> fPDGToName = {
          {2212, "Protons"},
          {211, "Pions"},
          {-11, "Electrons"},
          {-13, "Muons"},
          {321, "Kaons"},
          {-211, "Pions"},
          {11, "Electrons"},
          {13, "Muons"},
          {-321, "Kaons"}
        };

        std::string fSamplingFileName;
        std::string fResolutionFileName;
        TFile * fSamplingFile;
        TFile * fResolutionFile;
        std::map<std::string, THnSparseD *> fPDFs;
        std::map<std::string, TH2D *> fResolutionHists2D;
        void Setup1GeV(); //Shared by 2GeV
        void Setup3GeV();
        void Setup6GeV(); //Shared by 7GeV
        
        void Scale2DRes();
        void SetMinMax();

        std::string GetPrimaryEndProcess(const int &primary_pdg, const std::vector<int> & secondary_pdgs);
        std::string GetSecondaryProcess(const int &primary_pdg, const int &secondary_pdg);

        bool fSaveOutputTree;
        TTree * fOutputTree;
        TGraph * fTriggersGraph;
        int fOutputPDG;
        int fOutputEvent;
        double fOutputMomentum;
        double fOutputUnsmearedMomentum;
        double fOutputHUpstream, fOutputVUpstream;
        double fOutputHDownstream, fOutputVDownstream;

        bool fReduceNP04frontArea;
        
        // Number of beam interactions to overlay.
        int fOverlays;
        
        ifdh_ns::ifdh* fIFDH;
    };
}

// Create the random number generator
namespace {
  std::string const instanceName = "protoDUNEBeam";
}


//---------------------------------------------------------------------------------
//----------------------------------------constructors-----------------------------
evgen::ProtoDUNETriggeredBeam::ProtoDUNETriggeredBeam(fhicl::ParameterSet const & pset)
  : EDProducer{pset}
    // now create the engine (for example, use art); seed will be set
    // by calling declareEngine
  , fFlatRnd(createEngine(art::ServiceHandle<rndm::NuRandomService>{}->declareEngine(instanceName),
                          "HepJamesRandom", instanceName))
{
    // Call appropriate produces<>() functions here.
    produces< std::vector<simb::MCTruth> >();
    produces< sumdata::RunData, art::InRun >();
    produces< std::vector< beam::ProtoDUNEBeamEvent > >();
    // File reading variable initialisations
    fFileName = pset.get< std::string>("FileName");
    fBaseFileName = fFileName.substr(fFileName.rfind("/")+1);

    // Tree names
    fTOF1TreeName      = pset.get<std::string>("TOF1TreeName");
    fBPROF1TreeName    = pset.get<std::string>("BPROF1TreeName");
    fBPROF2TreeName    = pset.get<std::string>("BPROF2TreeName");
    fBPROF3TreeName    = pset.get<std::string>("BPROF3TreeName");
    fTRIG1TreeName     = pset.get<std::string>("TRIG1TreeName");
    fBPROFEXTTreeName  = pset.get<std::string>("BPROFEXTTreeName");
    fBPROF4TreeName    = pset.get<std::string>("BPROF4TreeName");
    fTRIG2TreeName      = pset.get<std::string>("TRIG2TreeName");
    fNP04frontTreeName = pset.get<std::string>("NP04frontTreeName");

    // Intensity variables
    fIntensity = pset.get<float>("Intensity");
    fReadoutWindow = pset.get<float>("ReadoutWindow");
    fBeamSpillLength = pset.get<float>("BeamSpillLength");

    fUseDataDriven = pset.get<bool>("UseDataDrivenPrimary");
    fReduceNP04frontArea = pset.get<bool>("ReduceNP04frontArea");
   
    // See if the user wants to start at an event other than zero.
    fStartEvent = pset.get<int>("StartEvent");
    
    // Or maybe there was --nskip specified in the command line or skipEvents in FHiCL?
    for (auto const & p : fhicl::ParameterSetRegistry::get())
    {
        if (p.second.has_key("source"))
        {
          // Need to add in another layer here because of some change (maybe in art?) 
          if (p.second.has_key("source.skipEvents")) 
          {
            fStartEvent += p.second.get<int>("source.skipEvents");
            break; // take the first occurence
          }
        } // no "else", if parameter not found, then just don't change anything
    }
    // ...and if there is -e option or firstEvent in FHiCL, this add up to the no. of events to skip.
    for (auto const & p : fhicl::ParameterSetRegistry::get())
    {
        if (p.second.has_key("source"))
        {
          if (p.second.has_key("source.firstEvent"))
          {
              int fe = p.second.get<int>("source.firstEvent") - 1; // events base index is 1
              if (fe > 0) fStartEvent += fe;
              break; // take the first occurence
          } // no "else", if parameter not found, then just don't change anything
        }
    }
    mf::LogInfo("ProtoDUNETriggeredBeam") << "Skip " << fStartEvent << " first events from the input file.";
    
    fEventNumber = fStartEvent;
    
    // Coordinate transform
    fBeamX = pset.get<float>("BeamX");
    fBeamY = pset.get<float>("BeamY");
    fBeamZ = pset.get<float>("BeamZ");
    fBeamThetaShift = pset.get<float>("BeamThetaShift",0.0);
    fBeamPhiShift   = pset.get<float>("BeamPhiShift",0.0);

    
    fRotateMonitorXZ = pset.get<float>("RotateMonitorXZ");
    fRotateMonitorYZ = pset.get<float>("RotateMonitorYZ");
    fBPROFEXTPos     = pset.get<float>("BPROFEXTPosZ");
    fBPROF4Pos       = pset.get<float>("BPROF4PosZ");
    fTRIG2Pos        = pset.get<float>("TRIG2PosZ");
    fNP04frontPos    = pset.get<float>("NP04frontPosZ");
    // Setup the beam monitor basis vectors
    BeamMonitorBasisVectors();   

    fIFDH = 0;
    
    fL1 = pset.get<float>("L1");
    fL2 = pset.get<float>("L2");
    fL3 = pset.get<float>("L3");
    fBeamBend = pset.get<float>("BeamBend");

    //New values for momentum spectrometer
    fLMag = pset.get<float>("LMag");
    fB    = pset.get<float>("B");
    fNominalP = pset.get<std::string>("NominalP");
    fLB = fB * fLMag * std::stod(fNominalP) / 7.;

    fMaxSamples = pset.get<int>("MaxSamples", 0);

    fRNG = TRandom3(pset.get<int>("Seed", 0));
    fVerbose = pset.get<bool>("Verbose", false);
    fIncludeAnti = pset.get<bool>("IncludeAnti", false);
    fResolutionFileName = pset.get<std::string>("ResolutionFileName");
    fSamplingFileName = pset.get<std::string>("SamplingFileName");
    
    fSaveOutputTree = pset.get<bool>("SaveOutputTree");

    // Make sure we use ifdh to open the beam input file.
    OpenInputFile(fFileName);
    if(fUseDataDriven){
      //std::string found_sampling_file = FindFile(fSamplingFileName);
      //OpenInputFile(found_sampling_file);

      //std::string found_res_file = FindFile(fResolutionFileName);
      //OpenInputFile(found_res_file);

      fSamplingFileName = FindFile(fSamplingFileName);
      OpenInputFile(fSamplingFileName);

      fResolutionFileName = FindFile(fResolutionFileName);
      OpenInputFile(fResolutionFileName);
    }
}



//-----------------------------default destructor----------------------------------
//------------------------------------------------------------------------------------
evgen::ProtoDUNETriggeredBeam::~ProtoDUNETriggeredBeam()
{
    fIFDH->cleanup();   
}

//-------------------------------------------------------------------------------------
void evgen::ProtoDUNETriggeredBeam::beginJob(){

    art::ServiceHandle<art::TFileService> tfs;
    
    TFile *inputFile = new TFile(fFileName.c_str(),"READ");
    // Check we have the file
    if(inputFile == 0x0){
        throw cet::exception("ProtoDUNETriggeredBeam") << "Input file " << fFileName << " cannot be read.\n";
    }
    
    TTree *frontFaceTree = (TTree*)inputFile->Get(fNP04frontTreeName.c_str());
    // Check we have the tree
    if(frontFaceTree == 0x0){
        throw cet::exception("ProtoDUNETriggeredBeam") << "Input tree " << fNP04frontTreeName << " cannot be read.\n";
    }
    std::cout << "All particle tree " << fNP04frontTreeName << " has " << frontFaceTree->GetEntries() << " entries" << std::endl;
   
    // Calculate the number of events to overlay
    CalculateNOverlays();

    // Fill all potential events from the NP04front tree
    FillParticleMaps(frontFaceTree);

    TTree *trig1Tree = (TTree*)inputFile->Get(fTRIG1TreeName.c_str());
    TTree *trig2Tree = (TTree*)inputFile->Get(fTRIG2TreeName.c_str());

    // Now search for trigger events
    std::vector<int> triggeredEventIDs =  FindTriggeredEvents(trig1Tree,trig2Tree);
    std::cout << "Proto trigger list has " << triggeredEventIDs.size() << " events" << std::endl; 

    // For triggered events, we now need to attach the other instrument information
    std::vector<std::string> otherInstrumentTreeNames;
    otherInstrumentTreeNames.push_back(fTOF1TreeName.c_str());
    otherInstrumentTreeNames.push_back(fBPROF1TreeName.c_str());
    otherInstrumentTreeNames.push_back(fBPROF2TreeName.c_str());
    otherInstrumentTreeNames.push_back(fBPROF3TreeName.c_str());
    otherInstrumentTreeNames.push_back(fBPROFEXTTreeName.c_str());
    otherInstrumentTreeNames.push_back(fBPROF4TreeName.c_str());

    for(const std::string treeName : otherInstrumentTreeNames){ 
      TTree *instrumentTree = (TTree*)inputFile->Get(treeName.c_str());
      FillInstrumentInformation(triggeredEventIDs,instrumentTree);
      std::cout << " - Finished adding information from " << treeName << std::endl;
    }
    std::cout << "Final trigger list has " << triggeredEventIDs.size() << " events" << std::endl; 
    fFinalTriggerEventIDs = triggeredEventIDs;

    // We are done with the input file now
    inputFile->Close();
    delete inputFile;
    inputFile = 0x0;

    // Data-driven file setup
    if(fUseDataDriven){
      fSamplingFile = new TFile(fSamplingFileName.c_str());
      fResolutionFile = new TFile(fResolutionFileName.c_str());
        if (fNominalP =="0.5") {
          // This is the same function as for 1 GeV
          Setup1GeV();
        }
        else if (fNominalP =="1") {
          Setup1GeV();
        }
        else if (fNominalP =="2") {
          // This is the same function as for 1 GeV
          Setup1GeV();
        }
        else if (fNominalP =="3") {
          Setup3GeV();
        }
        else if (fNominalP =="6") {
          Setup6GeV();
        }
        else if (fNominalP =="7") {
          // This is the same function as for 7 GeV
          Setup6GeV();
        }

      Scale2DRes();
      SetMinMax();
    }
    if (fSaveOutputTree) {
      fTriggersGraph = tfs->makeAndRegister<TGraph>("Triggers", fBaseFileName.c_str(), 1);
      fTriggersGraph->SetPoint(0, 0., triggeredEventIDs.size());
      if (fUseDataDriven) {
        fOutputTree = tfs->make<TTree>("tree", "");
        fOutputTree->Branch("PDG", &fOutputPDG);
        fOutputTree->Branch("Event", &fOutputEvent);
        fOutputTree->Branch("Momentum", &fOutputMomentum);
        fOutputTree->Branch("UnsmearedMomentum", &fOutputUnsmearedMomentum);
        fOutputTree->Branch("HUpstream", &fOutputHUpstream);
        fOutputTree->Branch("VUpstream", &fOutputVUpstream);
        fOutputTree->Branch("HDownstream", &fOutputHDownstream);
        fOutputTree->Branch("VDownstream", &fOutputVDownstream);
      }
    }
}

//----------------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::beginRun(art::Run& run)
{
    // Grab the geometry object to see what geometry we are using
    art::ServiceHandle<geo::Geometry> geo;
    std::unique_ptr<sumdata::RunData> runcol(new sumdata::RunData(geo->DetectorName()));
    run.put(std::move(runcol));
}

//--------------------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::endJob(){

}


//--------------------------------------------------------------------------------------------
void evgen::ProtoDUNETriggeredBeam::produce(art::Event & e)
{
    if(fEventNumber >= static_cast<int>(fFinalTriggerEventIDs.size())){
        throw cet::exception("ProtoDUNETriggeredBeam") << "Requested entry " << fEventNumber
        << " but tree only has entries 0 to "
        << fFinalTriggerEventIDs.size() - 1 << std::endl;
    }
   
    // Define the truth collection for this event.
    auto truthcol = std::make_unique< std::vector<simb::MCTruth> >();
    
    std::unique_ptr<std::vector<beam::ProtoDUNEBeamEvent> > beamData(new std::vector<beam::ProtoDUNEBeamEvent>);

    simb::MCTruth truth;
    beam::ProtoDUNEBeamEvent beamEvent;
 
    // Group the events together: a triggered event with fOverlay background events
    OverlaidTriggerEvent overlayEvent = GenerateOverlaidEvent(fFinalTriggerEventIDs.at(fEventNumber));

    // Fill the MCTruth object
    fOutputEvent = e.id().event();
    GenerateTrueEvent(truth, overlayEvent, beamEvent);
    
    // Add the MCTruth to the vector
    truthcol->push_back(truth);
    
    // Finally, add the MCTruth to the event
    e.put(std::move(truthcol));

    beamData->push_back( beamEvent );
    e.put( std::move( beamData ) );

    // We have made our event, increment the event number.
    ++fEventNumber;
}

//--------------------------------------------------------------------------------------

// Fill the particle maps using the input files. This links the events of interest
// to the entry number in fAllParticlesTree.
void evgen::ProtoDUNETriggeredBeam::FillParticleMaps(TTree *frontFaceTree){
    
    float eventID, trackID, pdgCode, parentID;
    float posX, posY, posZ, posT;
    float momX, momY, momZ;

    frontFaceTree->SetBranchAddress("EventID",&eventID);
    frontFaceTree->SetBranchAddress("TrackID",&trackID);
    frontFaceTree->SetBranchAddress("PDGid",&pdgCode);
    frontFaceTree->SetBranchAddress("ParentID",&parentID);
    frontFaceTree->SetBranchAddress("x",&posX);
    frontFaceTree->SetBranchAddress("y",&posY);
    frontFaceTree->SetBranchAddress("z",&posZ);
    frontFaceTree->SetBranchAddress("t",&posT);
    frontFaceTree->SetBranchAddress("Px",&momX);
    frontFaceTree->SetBranchAddress("Py",&momY);
    frontFaceTree->SetBranchAddress("Pz",&momZ);

    // Loop over all particles and group them by events
    for(unsigned int p = 0; p < frontFaceTree->GetEntries(); ++p){

      frontFaceTree->GetEntry(p);

      // Don't consider nuclei here
      const int intPdgCode = static_cast<int>(pdgCode);
      if(intPdgCode > 10000) continue;

      // If this particle is travelling backwards then it won't hit the detector
      if(momZ < 0) continue;

      // Convert to detector coordinate system
      ConvertCoordinates(posX,posY,posZ);

      // Keep only those particles that might reach the detector
      if(fReduceNP04frontArea){
        if(posX < -500 || posX > 500) continue;
        if(posY < -150 || posY > 850) continue;
      }

      // Convert momentum
      ConvertMomentum(momX,momY,momZ);

      const int intEventID = static_cast<int>(eventID);
      const int intTrackID = static_cast<int>(trackID);
      const int intParentID= static_cast<int>(parentID);
      BeamParticle newParticle(intTrackID, intPdgCode, intParentID,
                               posX, posY, posZ, posT, momX, momY, momZ);

      std::map<int,BeamEvent>::iterator evIter = fAllBeamEvents.find(intEventID);
      if(evIter == fAllBeamEvents.end()){
        BeamEvent newBeamEvent(intEventID);
        newBeamEvent.AddParticle(newParticle);
        fAllBeamEvents.insert(std::make_pair(intEventID,newBeamEvent));
      }
      else{
        evIter->second.AddParticle(newParticle); 
      }

    }

    std::cout << "Found " << fAllBeamEvents.size() << " potential events" << std::endl;

    // Reset the branch addresses as the variables are going out of scope
    frontFaceTree->ResetBranchAddresses();
}

//---------------------------------------------------------------------------------------

// Search for events that have particles in TRIG1 and TRIG2 planes
std::vector<int> evgen::ProtoDUNETriggeredBeam::FindTriggeredEvents(TTree *trig1Tree, TTree *trig2Tree){

  const std::vector<int> allowedPDGs = {11,-11,13,-13,211,-211,321,-321,2212};
  
  float eventID, trackID, pdgCode, parentID;
  float posX, posY, posZ, posT;
  float momX, momY, momZ;

  // Look at trigger two first to reduce computation
  trig2Tree->SetBranchAddress("EventID",&eventID);
  trig2Tree->SetBranchAddress("TrackID",&trackID);
  trig2Tree->SetBranchAddress("PDGid",&pdgCode);
  trig2Tree->SetBranchAddress("ParentID",&parentID);
  trig2Tree->SetBranchAddress("x",&posX);
  trig2Tree->SetBranchAddress("y",&posY);
  trig2Tree->SetBranchAddress("z",&posZ);
  trig2Tree->SetBranchAddress("t",&posT);
  trig2Tree->SetBranchAddress("Px",&momX);
  trig2Tree->SetBranchAddress("Py",&momY);
  trig2Tree->SetBranchAddress("Pz",&momZ);

  // Temporarily store the particle for events with particle in TRIG2. Just store
  // the first one if there are two
  std::map<int,BeamParticle> trig2Particles;

  for(unsigned int p = 0; p < trig2Tree->GetEntries(); ++p){
  
    trig2Tree->GetEntry(p);
  
    const int intEventID = static_cast<int>(eventID);

    // If this event didn't have any particles at NP04front then move on
    if(fAllBeamEvents.find(intEventID) == fAllBeamEvents.end()) continue;

    // Carry on if we already found a particle for this event
    if(trig2Particles.find(intEventID) != trig2Particles.end()) continue;

    // If the particle isn't of the type we want then move on
    if(std::find(allowedPDGs.begin(),allowedPDGs.end(),static_cast<int>(pdgCode))==allowedPDGs.end()) continue;
     

    TVector3 det_pos = ConvertProfCoordinates(posX, posY, posZ, fTRIG2Pos);
    BeamParticle particle(static_cast<int>(trackID), static_cast<int>(pdgCode), static_cast<int>(parentID),
                          det_pos.X(), det_pos.Y(), det_pos.Z(), posT, momX, momY, momZ);
                          //posX, posY, posZ, posT, momX, momY, momZ);
    trig2Particles.insert(std::make_pair(intEventID,particle));
  }

  trig2Tree->ResetBranchAddresses();

  // Now look at TRIG1
  trig1Tree->SetBranchAddress("EventID",&eventID);
  trig1Tree->SetBranchAddress("TrackID",&trackID);
  trig1Tree->SetBranchAddress("PDGid",&pdgCode);
  trig1Tree->SetBranchAddress("ParentID",&parentID);
  trig1Tree->SetBranchAddress("x",&posX);
  trig1Tree->SetBranchAddress("y",&posY);
  trig1Tree->SetBranchAddress("z",&posZ);
  trig1Tree->SetBranchAddress("t",&posT);
  trig1Tree->SetBranchAddress("Px",&momX);
  trig1Tree->SetBranchAddress("Py",&momY);
  trig1Tree->SetBranchAddress("Pz",&momZ);

  std::map<int,BeamParticle> trig1Particles;
  for(unsigned int p = 0; p < trig1Tree->GetEntries(); ++p){

    trig1Tree->GetEntry(p);
    const int intEventID = static_cast<int>(eventID);

    // Move on if this event had no particle in TRIG2
    if(trig2Particles.find(intEventID) == trig2Particles.end()) continue;

    if(std::find(allowedPDGs.begin(),allowedPDGs.end(),static_cast<int>(pdgCode))==allowedPDGs.end()) continue;
     
    BeamParticle particle(static_cast<int>(trackID), static_cast<int>(pdgCode), static_cast<int>(parentID),
                          posX, posY, posZ, posT, momX, momY, momZ);
    trig1Particles.insert(std::make_pair(intEventID,particle));
  }

  // Add the particle information at TRIG1 and TRIG2 to the triggered events
  const std::string trig1TreeName = fTRIG1TreeName.substr(fTRIG1TreeName.find("/")+1);
  const std::string trig2TreeName = fTRIG2TreeName.substr(fTRIG2TreeName.find("/")+1);

  std::vector<int> trigEventIDs;
  for(auto const &element : trig1Particles){
    BeamEvent &event = fAllBeamEvents.at(element.first);
    // Check that this makes sense... the same particle or one particle is the parent of the other
    if((element.second.fTrackID != trig2Particles.at(element.first).fTrackID) &&
       (element.second.fTrackID != trig2Particles.at(element.first).fParentID)) continue; 
    
    event.fTriggerID = trig2Particles.at(element.first).fTrackID;    
    event.fTriggeredParticleInfo.insert(std::make_pair(trig1TreeName,element.second));
    event.fTriggeredParticleInfo.insert(std::make_pair(trig2TreeName,trig2Particles.at(element.first)));

    bool isTriggerEvent = false;
    // There is a rare case where the TRIG2 particle can decay before NP04front
    if(event.fParticlesFront.find(event.fTriggerID) == event.fParticlesFront.end()){
      event.fHasInteracted = true;
      // Find the child particle in the map
      std::cout << "- Candidate event " << trigEventIDs.size() << " trigger particle of type " << trig2Particles.at(element.first).fPDG << " not at the front face... searching for children" << std::endl;
        for(const std::pair<int,BeamParticle> &partPair : event.fParticlesFront){
        if(partPair.second.fParentID == event.fTriggerID){
          std::cout << "  - Found child with PDG = " << partPair.second.fPDG << std::endl;
          event.fSecondaryTrackIDs.push_back(partPair.first);
          isTriggerEvent = true;
        }
      }
    }
    else{
      isTriggerEvent = true;
    }

    if(isTriggerEvent) trigEventIDs.push_back(element.first);
  }

  trig1Tree->ResetBranchAddresses();

  return trigEventIDs;
}

//---------------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::FillInstrumentInformation(std::vector<int> &eventIDs, TTree *instrumentTree){

  float eventID, trackID, pdgCode, parentID;
  float posX, posY, posZ, posT;
  float momX, momY, momZ;

  instrumentTree->SetBranchAddress("EventID",&eventID);
  instrumentTree->SetBranchAddress("TrackID",&trackID);
  instrumentTree->SetBranchAddress("PDGid",&pdgCode);
  instrumentTree->SetBranchAddress("ParentID",&parentID);
  instrumentTree->SetBranchAddress("x",&posX);
  instrumentTree->SetBranchAddress("y",&posY);
  instrumentTree->SetBranchAddress("z",&posZ);
  instrumentTree->SetBranchAddress("t",&posT);
  instrumentTree->SetBranchAddress("Px",&momX);
  instrumentTree->SetBranchAddress("Py",&momY);
  instrumentTree->SetBranchAddress("Pz",&momZ);

  // Buffer all of the tree entries for trigger events
  std::map<const int,std::vector<unsigned int>> triggerIndices; 
  std::map<const int,const bool> arePionDecays;
  std::map<const int,const int> trig1TrackIDs;
  std::map<const int,const int> trig2TrackIDs;
  std::map<const int,bool> foundTrackInEvent;
  for(const int &trigEventID : eventIDs){
    BeamEvent &event = fAllBeamEvents.at(trigEventID);
    const int trig1TrackID = event.fTriggeredParticleInfo.at("TRIG1").fTrackID;
    const int trig2TrackID = event.fTriggeredParticleInfo.at("TRIG2").fTrackID;
    const int trig1TrackPDG = event.fTriggeredParticleInfo.at("TRIG1").fPDG;
    const int trig2TrackPDG = event.fTriggeredParticleInfo.at("TRIG2").fPDG;
    const bool isPionDecay = ((trig1TrackPDG==211) && (trig2TrackPDG==-13)) || 
                             ((trig1TrackPDG==-211) && (trig2TrackPDG==13));

    arePionDecays.insert(std::make_pair(trigEventID,isPionDecay));
    trig1TrackIDs.insert(std::make_pair(trigEventID,trig1TrackID));
    trig2TrackIDs.insert(std::make_pair(trigEventID,trig2TrackID));
    foundTrackInEvent.insert(std::make_pair(trigEventID,false));

    triggerIndices.insert(std::make_pair(trigEventID,std::vector<unsigned int>()));
  }

  // Strip the first part of the tree name to get just the instrument name
  std::string treeName = instrumentTree->GetName();
  treeName = treeName.substr(treeName.find("/")+1);

  for(unsigned int p = 0; p < instrumentTree->GetEntries(); ++p){
    instrumentTree->GetEntry(p);
    // If this isn't a triggered event then move on
    const int thisEvent = static_cast<int>(eventID);
    if(std::find(eventIDs.begin(),eventIDs.end(),thisEvent)==eventIDs.end()) continue;

    triggerIndices.at(thisEvent).push_back(p);
    const int thisParticle = static_cast<int>(trackID);
    if(trig2TrackIDs.at(thisEvent) == thisParticle){
      BeamParticle particle(static_cast<int>(trackID), static_cast<int>(pdgCode), static_cast<int>(parentID),
                            posX, posY, posZ, posT, momX, momY, momZ);
      fAllBeamEvents.at(thisEvent).fTriggeredParticleInfo.insert(std::make_pair(treeName,particle));
      foundTrackInEvent.at(thisEvent) = true;
    }
  }


  for (auto it = eventIDs.begin(); it != eventIDs.end();) {
    const int ev = *it;
    if(foundTrackInEvent.at(ev)) {
      ++it;
      continue;
    }

    // If we didn't find TRIG2 particle, then look for the TRIG1 one
    for(const unsigned int index : triggerIndices.at(ev)){
      instrumentTree->GetEntry(index);
      const int thisEvent = static_cast<int>(eventID);
      const int thisParticle = static_cast<int>(trackID);
      if(trig1TrackIDs.at(thisEvent) == thisParticle){
        BeamParticle particle(static_cast<int>(trackID), static_cast<int>(pdgCode), static_cast<int>(parentID),
                            posX, posY, posZ, posT, momX, momY, momZ);
        fAllBeamEvents.at(thisEvent).fTriggeredParticleInfo.insert(std::make_pair(treeName,particle));
        foundTrackInEvent.at(thisEvent) = true;
        break;
      }
    }
  

    if(foundTrackInEvent.at(ev)) {
      ++it;
      continue;
    }
    // In the rare case that we still don't have the particle, try the TRIG1 parent
    const int parentTrack = fAllBeamEvents.at(ev).fTriggeredParticleInfo.at("TRIG1").fParentID;
    for(const unsigned int index : triggerIndices.at(ev)){
      instrumentTree->GetEntry(index);
      const int thisEvent = static_cast<int>(eventID);
      const int thisParticle = static_cast<int>(trackID);
      if(parentTrack == thisParticle){
        BeamParticle particle(static_cast<int>(trackID), static_cast<int>(pdgCode), static_cast<int>(parentID),
                            posX, posY, posZ, posT, momX, momY, momZ);
        fAllBeamEvents.at(thisEvent).fTriggeredParticleInfo.insert(std::make_pair(treeName,particle));
        foundTrackInEvent.at(thisEvent) = true;

        break;
      }

    }

    if(foundTrackInEvent.at(ev) == false){
      fAllBeamEvents.at(ev).fTriggerID = -999;
      // Remove this event from the input vector
      it = eventIDs.erase(it);
      std::cout << "Issue found with event " << ev << ". Removing it from the trigger list - " << eventIDs.size() << " remain" << std::endl;
      //std::cout << "We didn't find tracks " << trig2TrackIDs.at(ev) << " or " << trig1TrackIDs.at(ev) << " in " << treeName << std::endl;
      //std::cout << " - PDGs: 1 = " << fAllBeamEvents.at(ev).fTriggeredParticleInfo.at("TRIG1").fPDG
      //          << " and 2 = " << fAllBeamEvents.at(ev).fTriggeredParticleInfo.at("TRIG2").fPDG << std::endl;
      //for(const unsigned int index : triggerIndices.at(ev)){
      //  instrumentTree->GetEntry(index);
      //  std::cout << "- Choice = " << static_cast<int>(trackID) << " :: " << static_cast<int>(pdgCode) << std::endl;
      //}
    }
    else {
      ++it; 
    }
  }

  instrumentTree->ResetBranchAddresses();
}

//---------------------------------------------------------------------------------------

// Group the events to overlay a number of background events on each trigger event
evgen::ProtoDUNETriggeredBeam::OverlaidTriggerEvent evgen::ProtoDUNETriggeredBeam::GenerateOverlaidEvent(const int &trigEventID)
{
  OverlaidTriggerEvent newTriggerEvent(trigEventID);
  // Look to see if any of these neighbouring events had particles
  for(int overlayID = trigEventID - (fOverlays / 2); overlayID < trigEventID + (fOverlays/2); ++overlayID){
    if(fAllBeamEvents.find(overlayID) == fAllBeamEvents.end()) continue;
    
    newTriggerEvent.AddOverlay(overlayID);
  }

  return newTriggerEvent;
}

//-------------------------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::GenerateTrueEvent(simb::MCTruth &mcTruth, const OverlaidTriggerEvent &overlayEvent, beam::ProtoDUNEBeamEvent & beamEvent){
 
  // A single particle seems the most accurate description.
  mcTruth.SetOrigin(simb::kSingleParticle);

  // Get the actual triggered event first and the beam particle
  const BeamEvent trigEvent = fAllBeamEvents.at(overlayEvent.fTriggerEventID);

  // We need to be slightly careful here... there are rare events where the pion decays between TRIG2 and NP04front
  BeamParticle trigParticle;
  int primaryStatus = 1; // 1 means track in G4, 0 means don't track
  if(!trigEvent.fHasInteracted){
    trigParticle = trigEvent.fParticlesFront.at(trigEvent.fTriggerID);
  }
  else{
    const std::string trig2Name = fTRIG2TreeName.substr(fTRIG2TreeName.find("/")+1);
    trigParticle = trigEvent.fTriggeredParticleInfo.at(trig2Name);
    primaryStatus = 0;
  }

  std::cout << "- Generating event with trigger particle type " << trigParticle.fPDG << std::endl;

  // Time of the triggered particle (will make all times relative to this)
  const float triggerParticleTime = trigParticle.fPosT;

  // The track ID for primary particles in LArSoft should be negative. This is -1 for our triggered particle
  int trigOutputTrackID = -1*(mcTruth.NParticles() + 1);

  simb::MCParticle triggerParticle;
  if(!fUseDataDriven || trigEvent.fHasInteracted){
    // Create the MCParticle for the triggered beam particle - NB the time offset sets T = 0 for this particle
    triggerParticle = BeamParticleToMCParticle(trigParticle, trigOutputTrackID, triggerParticleTime, primaryStatus, "primary");

    // Fill the ProtoDUNEBeamEvent here to store beamline information
    SetBeamEvent(beamEvent,trigEvent);
    std::cout << "  - Created trigger particle using pure simulation" << std::endl;
    std::cout << "  - Active? " << primaryStatus << std::endl;
    std::cout << "  - Located at " << triggerParticle.EndX() << " " <<
                 triggerParticle.EndY() << " " << triggerParticle.EndZ() <<
                 std::endl;
  }
  else{
    triggerParticle = DataDrivenMCParticle(trigParticle, trigOutputTrackID, triggerParticleTime, beamEvent, primaryStatus, "primary");
    std::cout << "  - Created trigger particle using data driven method" << std::endl;
    std::cout << "  - Located at " << triggerParticle.EndX() << " " <<
                 triggerParticle.EndY() << " " << triggerParticle.EndZ() <<
                 std::endl;
  }

  //mcTruth.Add(triggerParticle);

  std::vector<simb::MCParticle> secondaries;
  std::vector<int> secondary_pdgs;
  // Now we can add any secondaries produced in the interaction before NP04front
  if(trigEvent.fHasInteracted){
    int count = mcTruth.NParticles();
    for(const int &id : trigEvent.fSecondaryTrackIDs){
      trigOutputTrackID = -1*(count + 2);
      BeamParticle secondary = trigEvent.fParticlesFront.at(id);
      simb::MCParticle secondaryParticle = BeamParticleToMCParticle(
          secondary, trigOutputTrackID, triggerParticleTime+secondary.fPosT, 1,
          GetSecondaryProcess(trigParticle.fPDG, secondary.fPDG), triggerParticle.TrackId());
      secondaries.push_back(secondaryParticle);
      triggerParticle.AddDaughter(trigOutputTrackID);
      secondary_pdgs.push_back(secondary.fPDG);
      //mcTruth.Add(secondaryParticle);
      ++count;
    }
  }

  if(!secondaries.empty()){
    std::cout << "  - Trigger particle has daughters:";
    for (int i = 0; i < triggerParticle.NumberDaughters(); ++i) {
      std::cout << " " << triggerParticle.Daughter(i);
    }
    std::cout << std::endl;
  }

  // Add the trigger particle now that the hierarchy has been established
  if (trigEvent.fHasInteracted) {
    triggerParticle.SetEndProcess(
        GetPrimaryEndProcess(trigParticle.fPDG, secondary_pdgs));
  }
  mcTruth.Add(triggerParticle);
  // Now add all of the secondaries (if any)
  for (const simb::MCParticle & sec : secondaries) {
    mcTruth.Add(sec);
    std::cout << "  - Added secondary " << sec.TrackId() <<
                 " of type " << sec.PdgCode() << " with process " <<
                 sec.Process() <<
                 " from interacting primary " << triggerParticle.PdgCode() <<
                 " " << triggerParticle.Mother() << std::endl; 
    std::cout << "  - Located at " << sec.EndX() << " " << 
                 sec.EndY() << " " <<
                 sec.EndZ() << std::endl;
  }

  // Now let's deal with all of the background events
  for(const int &eventID : overlayEvent.fOverlayEventIDs){
    // Each overlay event needs a base time within +/- fBeamWindow of the triggered beam event
    double baseTime = (fFlatRnd.fire() - 0.5)*2.0*(fReadoutWindow*1000.*1000.);
    // Special case for triggered event
    if(overlayEvent.fTriggerEventID == eventID) baseTime = triggerParticleTime;
  
    for (std::pair<int,BeamParticle> element : fAllBeamEvents.at(eventID).fParticlesFront){
      // Don't double count the trigger particle
      if(overlayEvent.fTriggerEventID == eventID && element.first == trigParticle.fTrackID) continue;
      BeamParticle particle = element.second;

      const int outputTrackID =  -1*(mcTruth.NParticles() + 1);
      // For background particles we need to "back-strapolate" them to BPROFEXT so that they can hit the CRT
      SetBackgroundPosition(particle);
      simb::MCParticle backgroundParticle = BeamParticleToMCParticle(particle,outputTrackID,baseTime,1,"primaryBackground");
      mcTruth.Add(backgroundParticle);
    }  

  }

  std::cout << "Created event with " << mcTruth.NParticles() << " particles." << std::endl;
}

//---------------------------------------------------------------------------------------

simb::MCParticle evgen::ProtoDUNETriggeredBeam::BeamParticleToMCParticle(const BeamParticle &beamParticle, const int outputTrackID, const float timeOffset, const int primaryStatus, const std::string process, const int motherID){

  simb::MCParticle newParticle(outputTrackID,beamParticle.fPDG,process, motherID, -1.0, primaryStatus);

  // Get the position four-vector
  const TLorentzVector pos(beamParticle.fPosX,beamParticle.fPosY,beamParticle.fPosZ,beamParticle.fPosT - timeOffset);

  // Get the mass to calculate the momentum four-vector
  const TDatabasePDG* databasePDG = TDatabasePDG::Instance();
  const TParticlePDG* definition = databasePDG->GetParticle(beamParticle.fPDG);
  const float mass = definition->Mass();
  const float energy = sqrt(mass*mass + TVector3(beamParticle.fMomX,beamParticle.fMomY,beamParticle.fMomZ).Mag2());

  const TLorentzVector mom(beamParticle.fMomX,beamParticle.fMomY,beamParticle.fMomZ,energy);

  // Add the single trajectory point to the MCParticle
  newParticle.AddTrajectoryPoint(pos,mom);
  return newParticle;
}

//---------------------------------------------------------------------------------------

simb::MCParticle evgen::ProtoDUNETriggeredBeam::DataDrivenMCParticle(
    const BeamParticle &beamParticle, const int outputTrackID,
    const float timeOffset, beam::ProtoDUNEBeamEvent & beamEvent, const int primaryStatus, const std::string process) {
  
  const int pdg = (fIncludeAnti ? beamParticle.fPDG : abs(beamParticle.fPDG));
  simb::MCParticle newParticle(outputTrackID, pdg, process, -1, -1.0, primaryStatus);
  
  double kin_samples[5]; //the point in phase space to check against pdf
  double pdf_check; //the number used for the checking
  bool sample_again = true;

  double sampled_momentum = 0.;
  double sampled_h_upstream = 0.;
  double sampled_v_upstream = 0.;
  double sampled_h_downstream = 0.;
  double sampled_v_downstream = 0.;
  int nSamples = 0;
  while (sample_again) {
    /*
    if (nSamples > fMaxSamples) {
      throw cet::exception("ProtoDUNETriggeredBeam") << 
          "Reached max samples. Exiting" << std::endl;
    }*/
    fRNG.RndmArray(5, &kin_samples[0]);
    pdf_check = fRNG.Rndm();

    //Need to convert the numbers sampled for the kinematics (0, 1)
    //to within the sampling range
    double kin_point[5];
    //Convert(kin_samples, minima, maxima, kin_point);
    ConvertSamplingPoint(kin_samples, fMinima, fMaxima, kin_point);


    //Find the bin in the THnSparseD. If the bin has a value of 0,
    //then it would not have been allocated to save on memory.
    //The false parameter prevents that bin from being allocated here,
    //to save memory
    const long long bin = fPDFs.at(fPDGToName.at(pdg))->GetBin(&kin_point[0],false);
    //The bin has no chance of being populated, move on
    if (bin == -1) continue;

    //Find how likely we are to populate this bin
    const double pdf_value = fPDFs.at(fPDGToName.at(pdg))->GetBinContent(bin);

    //If successful, save info and move on
    if (pdf_check <= pdf_value) {
      if (fVerbose) {
        std::cout << "bin: " << bin << " PDF val: " << pdf_value <<
                     " Check: " << pdf_check << std::endl;
        std::cout << kin_samples[0] << " " << kin_samples[1] << " " <<
                     kin_samples[2] << " " << kin_samples[3] << " " <<
                     kin_samples[4] << std::endl;
        std::cout << kin_point[0] << " " << kin_point[1] << " " <<
                     kin_point[2] << " " << kin_point[3] << " " <<
                     kin_point[4] << std::endl;
        std::cout << "Will keep" << std::endl;
      }

      //diff
      sampled_momentum = kin_point[0];
      sampled_v_upstream = kin_point[1];
      sampled_h_upstream = kin_point[2];
      sampled_v_downstream = kin_point[3];
      sampled_h_downstream = kin_point[4];

      sample_again = false;
    }
    ++nSamples;
  }

  TLorentzVector position(0, 0, 0, 0);
  TLorentzVector momentum(0., 0., 0., 0.);
  SetDataDrivenPosMom(position, momentum, sampled_h_upstream,
                      sampled_v_upstream, sampled_h_downstream,
                      sampled_v_downstream, sampled_momentum,
                      pdg);
  newParticle.AddTrajectoryPoint(position, momentum);

  SetDataDrivenBeamEvent(beamEvent, sampled_h_upstream,
                      sampled_v_upstream, sampled_h_downstream,
                      sampled_v_downstream, sampled_momentum);

  if (fSaveOutputTree) {
    fOutputPDG = pdg;
    fOutputMomentum = sampled_momentum;
    fOutputHUpstream = sampled_h_upstream;
    fOutputVUpstream = sampled_v_upstream;
    fOutputHDownstream = sampled_h_downstream;
    fOutputVDownstream = sampled_v_downstream;
    fOutputTree->Fill();
  }

  return newParticle;
}

//------------------------------------------------------------------------------------


void evgen::ProtoDUNETriggeredBeam::ConvertSamplingPoint(
    double input_point[5], std::vector<double> minima,
    std::vector<double> maxima, double output_point[5]) {
  for (int i = 0; i < 5; ++i) {
    const double delta = maxima[i] - minima[i];
    output_point[i] = minima[i] + delta * input_point[i];
  }
}

void evgen::ProtoDUNETriggeredBeam::SetDataDrivenPosMom(
    TLorentzVector & position, TLorentzVector & momentum, double sampledHUp,
    double sampledVUp, double sampledHDown, double sampledVDown,
    double sampledMomentum, int beamPDG) {

  const double upstreamX = 96. - sampledHUp; 
  const double upstreamY = 96. - sampledVUp;
  const double downstreamX = 96. - sampledHDown;
  const double downstreamY = 96. - sampledVDown;


//rename these
  TVector3 upstream_point = ConvertProfCoordinates(upstreamX, upstreamY, 0.,
                                                   fBPROFEXTPos);
  TVector3 downstream_point = ConvertProfCoordinates(downstreamX, downstreamY, 0.,
                                                     fBPROF4Pos);
  TVector3 dR = (downstream_point - upstream_point).Unit();

  //Project to generator point
  double deltaZ = (fBeamZ - downstream_point.Z());
  double deltaX = (dR.X() / dR.Z()) * deltaZ;
  double deltaY = (dR.Y() / dR.Z()) * deltaZ;

  TVector3 generator_point = downstream_point +
                             TVector3(deltaX, deltaY, deltaZ);
  //Set the position 4-vector
  //Time = 0 for now?
  position = TLorentzVector(generator_point, 0.);

  //Prints out the projected position at the face of the TPC
  if (fVerbose) {
    deltaZ = (-1.*fBeamZ);
    deltaX = (dR.X() / dR.Z()) * deltaZ;
    deltaY = (dR.Y() / dR.Z()) * deltaZ;

    TVector3 last_point = generator_point +
                          TVector3(deltaX, deltaY, deltaZ);

    std::cout << last_point.X() << " " << last_point.Y() << " " <<
                 last_point.Z() << std::endl;
  }

  double unsmeared_momentum = 0.;
  /*switch (fUnsmearType) {
    case 1:
      unsmeared_momentum = UnsmearMomentum1D(sampledMomentum, beamPDG);
      break;
    case 2:*/
  unsmeared_momentum = UnsmearMomentum2D(sampledMomentum, beamPDG);
      //GetSystWeights();
      /*break;
    default:
      //Just do 1D
      unsmeared_momentum = UnsmearMomentum1D(sampledMomentum, beamPDG);
      break;
  }*/

  //TVector3 mom_vec = fRandMomentum[fCurrentEvent]*dR;
  //TVector3 mom_vec = unsmeared_momentum*dR;
  const TVector3 mom_vec = unsmeared_momentum*dR;
  fOutputUnsmearedMomentum = unsmeared_momentum;

  //Get the PDG and set the mass & energy accordingly
  const TDatabasePDG * dbPDG = TDatabasePDG::Instance();
  const TParticlePDG * def = dbPDG->GetParticle(beamPDG);
  const double mass = def->Mass();
  const double energy = sqrt(mass*mass + mom_vec.Mag2());

  //Set the momentum 4-vector
  momentum = TLorentzVector(mom_vec, energy);
}

/*
double evgen::ProtoDUNETriggeredBeam::(double momentum, int pdg) {
  TF1 * res = fResolutions[fPDGToName[pdg]];
  double mean = res->GetParameter(1);
  double sigma = res->GetParameter(2);
  double t = fRNG.Gaus(mean, sigma); //random number from momentum resolution
  return (momentum/(t + 1.));
}*/

double evgen::ProtoDUNETriggeredBeam::UnsmearMomentum2D(double momentum, int pdg) {

  if (fVerbose) {
    std::cout << "Using 2D Unsmear method" << std::endl;
  }

  TH2D * res = fResolutionHists2D.at(fPDGToName.at(pdg));

  const int xBin = res->GetXaxis()->FindBin(momentum);
  if (fVerbose) {
    std::cout << "Momentum & bin: " << momentum << " " <<
                 xBin << std::endl;
  }

  double true_min = res->GetYaxis()->GetXmin();
  double true_max = res->GetYaxis()->GetXmax();
  for (int i = 1; i <= res->GetNbinsY(); ++i) {
    if (res->GetBinContent(xBin, i) > 0.) {
      true_min = res->GetYaxis()->GetBinLowEdge(i);
      break;
    }
  }
  for (int i = res->GetNbinsY(); i >= 1; --i) {
    if (res->GetBinContent(xBin, i) > 0.) {
      true_max = res->GetYaxis()->GetBinUpEdge(i);
      break;
    }
  }

  if (fVerbose)
    std::cout << "True min and max: " << true_min << " " << true_max << std::endl;
  
  double unsmeared_momentum = 0.;
  while (true) {
    const double t = fRNG.Uniform(true_min, true_max);
    const double pdf_check = fRNG.Rndm(); //random number to check against PDF 
    
    const int yBin = res->GetYaxis()->FindBin(t);
    const double pdf_value = res->GetBinContent(xBin, yBin);
    if (fVerbose) {
      std::cout << "True mom & bin: " << t << " " << yBin << std::endl;
      std::cout << "Check & val: " << pdf_check << " " << pdf_value <<
                   std::endl;
    }
    if (pdf_check < pdf_value) {
      unsmeared_momentum = t;
      if (fVerbose) std::cout << "Setting momentum to " << t << std::endl;
      break;
    }
  }

  return unsmeared_momentum;
}

std::string evgen::ProtoDUNETriggeredBeam::GetPrimaryEndProcess(
    const int &primary_pdg, const std::vector<int> &secondary_pdgs) {
  
  if (primary_pdg == 2212) {
    return "protonInelastic";
  }
  else if (primary_pdg == 211) {
    if (std::find(secondary_pdgs.begin(), secondary_pdgs.end(), -13) !=
        secondary_pdgs.end()) {
      return "Decay";
    }
    else {
      return "pi+Inelastic";
    }
  }
  else if (primary_pdg == -211) {
    if (std::find(secondary_pdgs.begin(), secondary_pdgs.end(), 13) !=
        secondary_pdgs.end()) {
      return "Decay";
    }
    else {
      return "pi-Inelastic";
    }
  }
  else if (primary_pdg == 321) {
    return "Decay";
  }

  return "default";
}

std::string evgen::ProtoDUNETriggeredBeam::GetSecondaryProcess(
    const int &primary_pdg, const int &secondary_pdg) {
  
  std::string preamble = "primary:";
  if (primary_pdg == 2212) {
    return preamble + "protonInelastic";
  }
  else if (primary_pdg == 211) {
    if (secondary_pdg == -13) {
      return preamble + "Decay";
    }
    else if (secondary_pdg == 211 || secondary_pdg == -211 ||
            secondary_pdg == 2212 || secondary_pdg == 2112 ||
            secondary_pdg > 2212 || secondary_pdg == 22) {
      return preamble + "pi+Inelastic";
    }
  }
  else if (primary_pdg == -211) {
    if (secondary_pdg == 13) {
      return preamble + "Decay";
    }
    else if (secondary_pdg == 211 || secondary_pdg == -211 ||
            secondary_pdg == 2212 || secondary_pdg == 2112 ||
            secondary_pdg > 2212 || secondary_pdg == 22) {
      return preamble + "pi-Inelastic";
    }
  }
  else if (primary_pdg == 321) {
    return preamble + "Decay";
  }

  std::cout << "Notice! Unknown secondary option " << primary_pdg << " " <<
               secondary_pdg << std::endl;
  return preamble + "default";
}

void evgen::ProtoDUNETriggeredBeam::SetMinMax() {
  if (fNominalP =="0.5") {
    fMinima[0] = 0.3;
    fMaxima[0] = 0.7;
  }
  else if (fNominalP =="2") {
    fMinima[0] = 1.6;
    fMaxima[0] = 2.4;
  }
  else if (fNominalP =="3") {
    fMinima[0] = 2.4;
    fMaxima[0] = 3.6;
  }
  else if (fNominalP =="6") {
    fMinima[0] = 5.0;
    fMaxima[0] = 7.0;
  }
  else if (fNominalP =="7") {
    fMinima[0] = 6.0;
    fMaxima[0] = 8.0;
  }
}

void evgen::ProtoDUNETriggeredBeam::Scale2DRes() {
  for (auto it = fResolutionHists2D.begin();
       it != fResolutionHists2D.end(); ++it) {
    TH2D * this_hist = it->second;
    for (int i = 1; i <= this_hist->GetNbinsX(); ++i) {
      const double integral = this_hist->TH1::Integral(i, i);
      double total = 0.;
      for (int j = 1; j <= this_hist->GetNbinsY(); ++j) {
        this_hist->SetBinContent(i, j,
            this_hist->GetBinContent(i, j) / integral);
        total += this_hist->GetBinContent(i, j);
      }
    }
  }

/*
  for (auto it = fResolutionHists2DPlus.begin();
       it != fResolutionHists2DPlus.end(); ++it) {
    TH2D * this_hist = it->second;
    for (int i = 1; i <= this_hist->GetNbinsX(); ++i) {
      double integral = this_hist->Integral(i, i);
      double total = 0.;
      for (int j = 1; j <= this_hist->GetNbinsY(); ++j) {
        this_hist->SetBinContent(i, j,
            this_hist->GetBinContent(i, j) / integral);
        total += this_hist->GetBinContent(i, j);
      }
    }

    this_hist->Divide(fResolutionHists2D[it->first]);
  }

  for (auto it = fResolutionHists2DMinus.begin();
       it != fResolutionHists2DMinus.end(); ++it) {
    TH2D * this_hist = it->second;
    for (int i = 1; i <= this_hist->GetNbinsX(); ++i) {
      double integral = this_hist->Integral(i, i);
      double total = 0.;
      for (int j = 1; j <= this_hist->GetNbinsY(); ++j) {
        this_hist->SetBinContent(i, j,
            this_hist->GetBinContent(i, j) / integral);
        total += this_hist->GetBinContent(i, j);
      }
    }

    this_hist->Divide(fResolutionHists2D[it->first]);
  }
  */
}

void evgen::ProtoDUNETriggeredBeam::Setup1GeV() {
  std::vector<std::string> particle_types = {
    "Muons", "Pions", "Protons", "Electrons"
  };
  for (size_t i = 0; i < particle_types.size(); ++i) {
    const std::string part_type = particle_types[i];
    if (part_type == "Muons") {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get("Pions"); 
    }
    else {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get(part_type.c_str());
    }

    //Also get the resolutions
    std::string res_name = "";
    if (part_type == "Muons") {
      res_name = "hPionsRes";
    }
    else {
      res_name = "h" + part_type + "Res";
    }

    //fResolutionHists[part_type] = (TH1D*)fResolutionFile->Get(res_name.c_str());

    res_name += "2D";
    fResolutionHists2D[part_type] = (TH2D*)fResolutionFile->Get(res_name.c_str());

    /*
    std::string plus_name = res_name + "Plus";
    fResolutionHists2DPlus[part_type] = (TH2D*)fResolutionFile->Get(plus_name.c_str());

    std::string minus_name = res_name + "Minus";
    fResolutionHists2DMinus[part_type] = (TH2D*)fResolutionFile->Get(minus_name.c_str());
    */

  }
}

void evgen::ProtoDUNETriggeredBeam::Setup3GeV() {
  std::vector<std::string> particle_types = {
    "Muons", "Pions", "Protons", "Electrons", "Kaons"
  };

  for (size_t i = 0; i < particle_types.size(); ++i) {
    const std::string part_type = particle_types[i];
    if (part_type == "Muons") {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get("Pions"); 
    }
    else if (part_type == "Kaons") {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get("Protons"); 
    }
    else {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get(part_type.c_str());
    }

    //Also get the resolutions
    std::string res_name = "";
    if (part_type == "Muons") {
      res_name = "hPionsRes";
    }
    /*
    else if (part_type == "Kaons") {
      res_name = "hProtonsRes";
    }*/
    else {
      res_name = "h" + part_type + "Res";
    }
    
    //fResolutionHists[part_type] = (TH1D*)fResolutionFile->Get(res_name.c_str());

    res_name += "2D";
    fResolutionHists2D[part_type] = (TH2D*)fResolutionFile->Get(res_name.c_str());

    /*
    std::string plus_name = res_name + "Plus";
    fResolutionHists2DPlus[part_type] = (TH2D*)fResolutionFile->Get(plus_name.c_str());

    std::string minus_name = res_name + "Minus";
    fResolutionHists2DMinus[part_type] = (TH2D*)fResolutionFile->Get(minus_name.c_str());
    */
  }
}

void evgen::ProtoDUNETriggeredBeam::Setup6GeV() {
  std::vector<std::string> particle_types = {
    "Muons", "Pions", "Protons", "Electrons", "Kaons"
  };

  for (size_t i = 0; i < particle_types.size(); ++i) {
    const std::string part_type = particle_types[i];
    if (part_type == "Muons" || part_type == "Electrons") {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get("Pions"); 
    }
    else {
      fPDFs[part_type] = (THnSparseD*)fSamplingFile->Get(part_type.c_str());
    }

    //Also get the resolutions
    std::string res_name = "";
    if (part_type == "Muons") {
      res_name = "hPionsRes";
    }
    else {
      res_name = "h" + part_type + "Res";
    }

    //fResolutionHists[part_type] = (TH1D*)fResolutionFile->Get(res_name.c_str());

    res_name += "2D";
    fResolutionHists2D[part_type] = (TH2D*)fResolutionFile->Get(res_name.c_str());

    /*
    std::string plus_name = res_name + "Plus";
    fResolutionHists2DPlus[part_type] = (TH2D*)fResolutionFile->Get(plus_name.c_str());

    std::string minus_name = res_name + "Minus";
    fResolutionHists2DMinus[part_type] = (TH2D*)fResolutionFile->Get(minus_name.c_str());
    */
  }
}

void evgen::ProtoDUNETriggeredBeam::SetDataDrivenBeamEvent(
    beam::ProtoDUNEBeamEvent & beamEvent,
    double sampledHUp, double sampledVUp, double sampledHDown,
    double sampledVDown, double sampledMomentum) {

  beamEvent.SetTOFs(std::vector<double>{0.});
  beamEvent.SetTOFChans(std::vector<int>{0});
  beamEvent.SetUpstreamTriggers(std::vector<size_t>{0});
  beamEvent.SetDownstreamTriggers(std::vector<size_t>{0});
  beamEvent.SetCalibrations(0., 0., 0., 0.);
  beamEvent.DecodeTOF();

  beamEvent.SetMagnetCurrent(0.);
  beamEvent.SetTimingTrigger(12);

  beam::CKov dummy;
  dummy.trigger = 0;
  dummy.pressure = 0.;
  dummy.timeStamp = 0.;
  beamEvent.SetCKov0(dummy);
  beamEvent.SetCKov1(dummy);

  beamEvent.SetActiveTrigger(0);
  beamEvent.SetT0(std::make_pair(0.,0.));

  //Dummy positions for these
  beamEvent.SetFBMTrigger("XBPF022697", MakeFiberMonitor(.5));
  beamEvent.SetFBMTrigger("XBPF022698", MakeFiberMonitor(.5));
  beamEvent.SetFBMTrigger("XBPF022701", MakeFiberMonitor(.5));
  beamEvent.SetFBMTrigger("XBPF022702", MakeFiberMonitor(.5));

  const double upstream_x = sampledHUp;
  const double upstream_y = sampledVUp;
  const double downstream_x = sampledHDown;
  const double downstream_y = sampledVDown;

  beamEvent.SetFBMTrigger("XBPF022707", MakeFiberMonitor(96. - upstream_x));//X
  beamEvent.SetFBMTrigger("XBPF022708", MakeFiberMonitor(96. - upstream_y));//Y

  beamEvent.SetFBMTrigger("XBPF022716", MakeFiberMonitor(96. - downstream_x));//X
  beamEvent.SetFBMTrigger("XBPF022717", MakeFiberMonitor(96. - downstream_y));//Y

  MakeTracks(beamEvent);

  beamEvent.AddRecoBeamMomentum(sampledMomentum);
}

// Function written in similar way as "openDBs()" in CORSIKAGen_module.cc
void evgen::ProtoDUNETriggeredBeam::OpenInputFile(std::string & filename)
{
    // Setup ifdh object
    if (!fIFDH)
    {
        fIFDH = new ifdh_ns::ifdh;
    }
    
    const char* ifdh_debug_env = std::getenv("IFDH_DEBUG_LEVEL");
    if ( ifdh_debug_env )
    {
        mf::LogInfo("ProtoDUNETriggeredBeam") << "IFDH_DEBUG_LEVEL: " << ifdh_debug_env<<"\n";
        fIFDH->set_debug(ifdh_debug_env);
    }
    
    const std::string path(gSystem->DirName(filename.c_str()));
    const std::string pattern(gSystem->BaseName(filename.c_str()));
    
    auto const flist = fIFDH->findMatchingFiles(path,pattern);
    if (flist.empty())
    {
        struct stat buffer;
        if (stat(filename.c_str(), &buffer) != 0)
        {
            throw cet::exception("ProtoDUNETriggeredBeam") << "No files returned for path:pattern: "<<path<<":"<<pattern<<std::endl;
        }
        else
        {
            mf::LogInfo("ProtoDUNETriggeredBeam") << "For "<< filename <<"\n";
        }
    }
    else
    {
        std::pair<std::string, long> f = flist.front();
        
        mf::LogInfo("ProtoDUNETriggeredBeam") << "For "<< filename <<"\n";
        
        // Do the fetching, store local filepaths in locallist
        
        mf::LogInfo("ProtoDUNETriggeredBeam")
        << "Fetching: " << f.first << " " << f.second <<"\n";
        std::string fetchedfile(fIFDH->fetchInput(f.first));
        MF_LOG_DEBUG("ProtoDUNETriggeredBeam") << " Fetched; local path: " << fetchedfile;
        
        filename = fetchedfile;
    }
}


//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::ConvertCoordinates(float &x, float &y, float &z){
    
    // Convert to cm and shift to the detector coordinate frame
    x = (x/10.) + fBeamX;
    y = (y/10.) + fBeamY;
    z = fBeamZ; // Just use the z position    
}

//--------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::ConvertMomentum(float &px, float &py, float &pz){
    
    // Convert to GeV
    px = px / 1000.;
    py = py / 1000.;
    pz = pz / 1000.;
   
    // If we want to rotate by changing theta and phi, do it here.
    TVector3 momVec(px,py,pz);    
    momVec.SetTheta(momVec.Theta() + fBeamThetaShift);
    momVec.SetPhi(momVec.Phi() + fBeamPhiShift);

    px = momVec.X();
    py = momVec.Y();
    pz = momVec.Z();
}

//-----------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::CalculateNOverlays(){
    
    // The number of events to overlay is as follows:
    // N = Intensity * 2.0 * ReadoutWindow / BeamSpillLength
    fOverlays = fIntensity * (2.0 * fReadoutWindow / 1000.) / fBeamSpillLength;
    std::cout << "Number of overlays = " << fOverlays << std::endl;   
}
 
//----------------------------------------------------------------------------------

// We need to rotate the beam monitor coordinates into the detector frame (matching NP04front)
// This means they can later be treated in the same way as the standard NP04front positions
TLorentzVector evgen::ProtoDUNETriggeredBeam::ConvertBeamMonitorCoordinates(float x, float y, float z, float t, float zOffset){

  const float off = fNP04frontPos - zOffset;

//  TLorentzVector old(x,y,z,t);

  // Convert the coordinates using the rotated basis vectors
  float newX = x*fBMBasisX.X() + y*fBMBasisY.X() + (z-zOffset)*fBMBasisZ.X() + off*fabs(fBMBasisZ.X());
  float newY = x*fBMBasisX.Y() + y*fBMBasisY.Y() + (z-zOffset)*fBMBasisZ.Y() + off*fabs(fBMBasisZ.Y());
  float newZ = x*fBMBasisX.Z() + y*fBMBasisY.Z() + (z-zOffset) - off*fabs(fBMBasisZ.Z());

  // Account for the small differences between NP04front and the detector coordinates
  newX += fBeamX*10.;
  newY += fBeamY*10.;
  newZ += fBeamZ*10.;

  // Make our new beam monitor position in the detector coordinate system
  TLorentzVector result(newX,newY,newZ,t);

//  std::cout << "Coordinate transform..." << std::endl;
//  old.Print();
//  result.Print();

  return result;
}
 
//----------------------------------------------------------------------------------

TVector3 evgen::ProtoDUNETriggeredBeam::ConvertProfCoordinates(double x, double y, double z, double zOffset){
  const double off = fNP04frontPos - zOffset;

//  TVector3 old(x,y,z);

  double newX = x*fBMBasisX.X() + y*fBMBasisY.X() + /*(z-zOffset)*fBMBasisZ.X()*/ + off*fabs(fBMBasisZ.X());
  double newY = x*fBMBasisX.Y() + y*fBMBasisY.Y() + /*(z-zOffset)*fBMBasisZ.Y()*/ + off*fabs(fBMBasisZ.Y());
  double newZ = x*fBMBasisX.Z() + y*fBMBasisY.Z() + /*(z-zOffset)              */ - off*fabs(fBMBasisZ.Z());

  newX += fBeamX*10.;
  newY += fBeamY*10.;
  newZ += fBeamZ*10.;

  TVector3 result(newX/10., newY/10., newZ/10.);
  return result;
}
 
//----------------------------------------------------------------------------------

TVector3 evgen::ProtoDUNETriggeredBeam::ConvertBeamMonitorMomentumVec(float px, float py, float pz){

  TVector3 newMom(px,py,pz);
  RotateMonitorVector(newMom);
  return newMom;
}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::BeamMonitorBasisVectors(){

  fBMBasisX = TVector3(1.,0.,0.);
  fBMBasisY = TVector3(0.,1.,0.);
  fBMBasisZ = TVector3(0.,0.,1.);
  RotateMonitorVector(fBMBasisX);
  RotateMonitorVector(fBMBasisY);
  RotateMonitorVector(fBMBasisZ);

}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::RotateMonitorVector(TVector3 &vec){

  // Note: reordering how these are done in order to keep the basis
  //       vectors of the monitors parallel to the ground. 
  vec.RotateX( fRotateMonitorYZ * TMath::Pi() / 180. );
  vec.RotateY( fRotateMonitorXZ * TMath::Pi() / 180. );

}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::SetBackgroundPosition(BeamParticle &particle){

  const TVector3 pos(particle.fPosX,particle.fPosY,particle.fPosZ);
  const TVector3 dir = TVector3(particle.fMomX,particle.fMomY,particle.fMomZ).Unit();

  // Want to move the position upstream by a distance equal to fNP04frontPos - fBPROFEXTPos
  // This length is in the beam direction frame unless we account for it
  // Convert the instrument positions from mm to cm
  const float shiftLength = (fNP04frontPos - fBPROFEXTPos)/(10.0*fBMBasisZ.Z());

  const TVector3 shiftedPos = pos - shiftLength*dir;

  particle.fPosX = shiftedPos.X();
  particle.fPosY = shiftedPos.Y();
  particle.fPosZ = shiftedPos.Z();
}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::SetBeamEvent(beam::ProtoDUNEBeamEvent & beamevt, const BeamEvent &triggerEvent){

  // Instrument names
  const std::string tof1Name     = fTOF1TreeName.substr(fTOF1TreeName.find("/")+1);
  const std::string bprof1Name   = fBPROF1TreeName.substr(fBPROF1TreeName.find("/")+1);
  const std::string bprof2Name   = fBPROF2TreeName.substr(fBPROF2TreeName.find("/")+1);
  const std::string bprof3Name   = fBPROF3TreeName.substr(fBPROF3TreeName.find("/")+1);
  const std::string trig1Name    = fTRIG1TreeName.substr(fTRIG1TreeName.find("/")+1);
  const std::string bprofEXTName = fBPROFEXTTreeName.substr(fBPROFEXTTreeName.find("/")+1);
  const std::string bprof4Name   = fBPROF4TreeName.substr(fBPROF4TreeName.find("/")+1);
  const std::string trig2Name    = fTRIG2TreeName.substr(fTRIG2TreeName.find("/")+1);

  // TOF first
  const float trig2Time = triggerEvent.fTriggeredParticleInfo.at(trig2Name).fPosT;
  const float tof1Time  = triggerEvent.fTriggeredParticleInfo.at(tof1Name).fPosT;
  beamevt.SetTOFs(std::vector<double>{trig2Time - tof1Time});
  beamevt.SetTOFChans( std::vector<int>{ 0 } );
  beamevt.SetUpstreamTriggers( std::vector<size_t>{0} );
  beamevt.SetDownstreamTriggers( std::vector<size_t>{0} );
  beamevt.SetCalibrations( 0., 0., 0., 0. );
  beamevt.DecodeTOF();

  // Fibre monitors
  const BeamParticle &bprof1Particle = triggerEvent.fTriggeredParticleInfo.at(bprof1Name);
  const BeamParticle &bprof2Particle = triggerEvent.fTriggeredParticleInfo.at(bprof2Name);
  const BeamParticle &bprof3Particle = triggerEvent.fTriggeredParticleInfo.at(bprof3Name);
  const BeamParticle &bprofExtParticle = triggerEvent.fTriggeredParticleInfo.at(bprofEXTName);
  const BeamParticle &bprof4Particle = triggerEvent.fTriggeredParticleInfo.at(bprof4Name);
  // (x,y) for BPROF1
  beamevt.SetFBMTrigger( "XBPF022697", MakeFiberMonitor( bprof1Particle.fPosX ) );
  beamevt.SetFBMTrigger( "XBPF022698", MakeFiberMonitor( bprof1Particle.fPosY ) );
  // Just x for BPROF2 and BPROF3
  beamevt.SetFBMTrigger( "XBPF022701", MakeFiberMonitor( bprof2Particle.fPosX ) );
  beamevt.SetFBMTrigger( "XBPF022702", MakeFiberMonitor( bprof3Particle.fPosX ) );
  // (x,y) for BPROFEXT
  beamevt.SetFBMTrigger( "XBPF022707", MakeFiberMonitor( bprofExtParticle.fPosX ) );
  beamevt.SetFBMTrigger( "XBPF022708", MakeFiberMonitor( bprofExtParticle.fPosY ) );
  // (x,y) for BPROF4
  beamevt.SetFBMTrigger( "XBPF022716", MakeFiberMonitor( bprof4Particle.fPosX ) );
  beamevt.SetFBMTrigger( "XBPF022717", MakeFiberMonitor( bprof4Particle.fPosY ) );

  // Cherenkovs aren't simulated, so set to dummy values
  beam::CKov dummy;
  dummy.trigger = 0;
  dummy.pressure = 0.;
  dummy.timeStamp = 0.;
  beamevt.SetCKov0( dummy );
  beamevt.SetCKov1( dummy );

  // Trigger information
  beamevt.SetMagnetCurrent( 0. );
  beamevt.SetTimingTrigger( 12 );
  beamevt.SetActiveTrigger(0);
  beamevt.SetT0( std::make_pair(0.,0.) );
 
  // Do the beamline instrumentation reconstruction
  MakeTracks( beamevt );
  MomentumSpectrometer( beamevt );

/*  
  //This will just use the class members
  beamevt.SetTOFs( std::vector<double>{ fGoodTRIG2_t - fGoodTOF1_t } );
  beamevt.SetTOFChans( std::vector<int>{ 0 } );
  beamevt.SetUpstreamTriggers( std::vector<size_t>{0} );
  beamevt.SetDownstreamTriggers( std::vector<size_t>{0} );
  beamevt.SetCalibrations( 0., 0., 0., 0. );
  beamevt.DecodeTOF();

  beamevt.SetMagnetCurrent( 0. );
  beamevt.SetTimingTrigger( 12 );

  beam::CKov dummy;
  dummy.trigger = 0;
  dummy.pressure = 0.;
  dummy.timeStamp = 0.;
  beamevt.SetCKov0( dummy );
  beamevt.SetCKov1( dummy );

  beamevt.SetActiveTrigger(0);
  beamevt.SetT0( std::make_pair(0.,0.) );

  beamevt.SetFBMTrigger( "XBPF022697", MakeFiberMonitor( fGoodBPROF1_x ) );
  beamevt.SetFBMTrigger( "XBPF022698", MakeFiberMonitor( fGoodBPROF1_y ) );
  beamevt.SetFBMTrigger( "XBPF022701", MakeFiberMonitor( fGoodBPROF2_x ) );
  beamevt.SetFBMTrigger( "XBPF022702", MakeFiberMonitor( fGoodBPROF3_x ) );

  beamevt.SetFBMTrigger( "XBPF022707", MakeFiberMonitor( fGoodBPROFEXT_x ) );
  beamevt.SetFBMTrigger( "XBPF022708", MakeFiberMonitor( fGoodBPROFEXT_y ) );

  beamevt.SetFBMTrigger( "XBPF022716", MakeFiberMonitor( fGoodBPROF4_x ) );
  beamevt.SetFBMTrigger( "XBPF022717", MakeFiberMonitor( fGoodBPROF4_y ) );

  MakeTracks( beamevt );
  MomentumSpectrometer( beamevt );


  if( fSaveOutputTree ){ 
    fReco_p = beamevt.GetRecoBeamMomentum(0);
    fReco_tof = beamevt.GetTOF();
    std::cout << "TOF: " << beamevt.GetTOFs()[0] << " " << beamevt.GetTOF() << std::endl;
    fNP04_PDG = fGoodNP04front_PDGid;

    fNP04front_p = sqrt( fGoodNP04front_Px * fGoodNP04front_Px 
                       + fGoodNP04front_Py * fGoodNP04front_Py 
                       + fGoodNP04front_Pz * fGoodNP04front_Pz );

    fXBPF697_p = sqrt( fGoodBPROF1_Px * fGoodBPROF1_Px 
                     + fGoodBPROF1_Py * fGoodBPROF1_Py 
                     + fGoodBPROF1_Pz * fGoodBPROF1_Pz );

    fXBPF701_p = sqrt( fGoodBPROF2_Px*fGoodBPROF2_Px 
                     + fGoodBPROF2_Py*fGoodBPROF2_Py 
                     + fGoodBPROF2_Pz*fGoodBPROF2_Pz );

    fXBPF702_p = sqrt( fGoodBPROF3_Px*fGoodBPROF3_Px 
                     + fGoodBPROF3_Py*fGoodBPROF3_Py 
                     + fGoodBPROF3_Pz*fGoodBPROF3_Pz );
    fXBPF697_x = fGoodBPROF1_x;
    fXBPF701_x = fGoodBPROF2_x;
    fXBPF702_x = fGoodBPROF3_x;

    fXBPF716_x = fGoodBPROF4_x;
    fXBPF717_y = fGoodBPROF4_y;
    fXBPF707_x = fGoodBPROFEXT_x;
    fXBPF708_y = fGoodBPROFEXT_y;

    
    fXBPF697_f = beamevt.GetFBM( "XBPF022697" ).active[0];
    fXBPF701_f = beamevt.GetFBM( "XBPF022701" ).active[0];
    fXBPF702_f = beamevt.GetFBM( "XBPF022702" ).active[0];

    fXBPF707_f = beamevt.GetFBM( "XBPF022707" ).active[0];
    fXBPF708_f = beamevt.GetFBM( "XBPF022708" ).active[0];
    fXBPF716_f = beamevt.GetFBM( "XBPF022716" ).active[0];
    fXBPF717_f = beamevt.GetFBM( "XBPF022717" ).active[0];

    fXBPF697_rx = GetPosition( fXBPF697_f ); 
    fXBPF701_rx = GetPosition( fXBPF701_f ); 
    fXBPF702_rx = GetPosition( fXBPF702_f ); 
                                         
    fXBPF716_rx = GetPosition( fXBPF716_f ); 
    fXBPF717_ry = GetPosition( fXBPF717_f ); 
    fXBPF707_rx = GetPosition( fXBPF707_f ); 
    fXBPF708_ry = GetPosition( fXBPF708_f ); 

    //fTrueFront_x = fGoodNP04front_x + fBeamX;
    //fTrueFront_y = fGoodNP04front_y + fBeamY;
    //fTrueFront_z = fGoodNP04front_z + fBeamZ;

    fRecoFront_x = beamevt.GetBeamTrack(0).End().X();
    fRecoFront_y = beamevt.GetBeamTrack(0).End().Y();
    fRecoFront_z = beamevt.GetBeamTrack(0).End().Z();

    fRecoTree->Fill();
  }
*/
}
 
//----------------------------------------------------------------------------------

beam::FBM evgen::ProtoDUNETriggeredBeam::MakeFiberMonitor( float pos ){
  beam::FBM theFBM;

  //I should probably just make this into
  //a constructor for the FBM...
  theFBM.ID = -1;
  theFBM.glitch_mask = {};
  std::uninitialized_fill( std::begin(theFBM.fiberData), std::end(theFBM.fiberData), 0. );
  std::uninitialized_fill( std::begin(theFBM.timeData), std::end(theFBM.timeData), 0. );
  theFBM.timeStamp = 0.;

  const short f = 96 -  short( floor(pos) ) - 1;
  theFBM.fibers[f] = 1;
  theFBM.active.push_back(f);
  theFBM.decoded = true;

  return theFBM; 
}
 
//----------------------------------------------------------------------------------

double evgen::ProtoDUNETriggeredBeam::GetPosition( short fiber ){
  return ((96 - fiber) - .5);
}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::MakeTracks( beam::ProtoDUNEBeamEvent & beamEvent ){
  
  //We should only have one active fiber at a time
  //
  //Might need to ask Leigh, etc. if it's possible
  //to have multiple particles going through at the
  //same time. In which case -- try to implement it
  
  const short fx1 = beamEvent.GetFBM( "XBPF022707" ).active[0];
  const short fy1 = beamEvent.GetFBM( "XBPF022708" ).active[0];

  const double x1 = GetPosition( fx1 );
  const double y1 = GetPosition( fy1 );

  TVector3 pos1 = ConvertProfCoordinates( x1, y1, 0., fBPROFEXTPos );

  const short fx2 = beamEvent.GetFBM( "XBPF022716" ).active[0];
  const short fy2 = beamEvent.GetFBM( "XBPF022717" ).active[0];

  const double x2 = GetPosition( fx2 );
  const double y2 = GetPosition( fy2 );

  TVector3 pos2 = ConvertProfCoordinates( x2, y2, 0., fBPROF4Pos );
 
  std::vector< TVector3 > thePoints = { pos1, pos2, ProjectToTPC( pos1, pos2 ) };
  std::vector< TVector3 > theMomenta = {
    ( pos2 - pos1 ).Unit(),
    ( pos2 - pos1 ).Unit(),
    ( pos2 - pos1 ).Unit()
  };

  beamEvent.AddBeamTrack(
    recob::Track(
      recob::TrackTrajectory(recob::tracking::convertCollToPoint( thePoints ),
                             recob::tracking::convertCollToVector( theMomenta ),
                             recob::Track::Flags_t( thePoints.size() ),
                             false ),
      0, -1., 0, recob::tracking::SMatrixSym55(), recob::tracking::SMatrixSym55(), 1 
    )
  );
    
}
 
//----------------------------------------------------------------------------------

TVector3 evgen::ProtoDUNETriggeredBeam::ProjectToTPC(TVector3 firstPoint, TVector3 secondPoint){
  const TVector3 dR = (secondPoint - firstPoint);
  
  const double deltaZ = -1.*secondPoint.Z();
  const double deltaX = deltaZ * (dR.X() / dR.Z());
  const double deltaY = deltaZ * (dR.Y() / dR.Z());

  TVector3 lastPoint = secondPoint + TVector3(deltaX, deltaY, deltaZ);
  return lastPoint;
}
 
//----------------------------------------------------------------------------------

void evgen::ProtoDUNETriggeredBeam::MomentumSpectrometer( beam::ProtoDUNEBeamEvent & beamEvent ){

  const short f1 = beamEvent.GetFBM( "XBPF022697" ).active[0];
  const short f2 = beamEvent.GetFBM( "XBPF022701" ).active[0];
  const short f3 = beamEvent.GetFBM( "XBPF022702" ).active[0];

  const double x1 = -1.e-3 * GetPosition( f1 );
  const double x2 = -1.e-3 * GetPosition( f2 );
  const double x3 = -1.e-3 * GetPosition( f3 );

  const double cos_theta = MomentumCosTheta( x1, x2, x3 );
  const double momentum = 299792458*fLB/(1.E9 * acos(cos_theta));
  beamEvent.AddRecoBeamMomentum( momentum );

}
 
//----------------------------------------------------------------------------------

double evgen::ProtoDUNETriggeredBeam::MomentumCosTheta(double x1, double x2, double x3){
  const double a =  (x2*fL3 - x3*fL2)*cos(fBeamBend)/(fL3-fL2);
 
  const double numTerm = (a - x1)*( (fL3 - fL2)*tan(fBeamBend) + (x3 - x2)*cos(fBeamBend) ) + fL1*( fL3 - fL2 );

  const double denomTerm1 = sqrt( fL1*fL1 + (a - x1)*(a - x1) );
  const double denomTerm2 = sqrt( TMath::Power( ( (fL3 - fL2)*tan(fBeamBend) + (x3 - x2)*cos(fBeamBend) ),2)
                   + TMath::Power( ( (fL3 - fL2) ),2) );
  const double denom = denomTerm1 * denomTerm2;

  const double cosTheta = numTerm/denom;  
  return cosTheta;
}


std::string evgen::ProtoDUNETriggeredBeam::FindFile(const std::string filename) {
  mf::LogInfo("evgen::ProtoDUNETriggeredBeam::FindFile") << "Searching for " << filename;
  if (cet::file_exists(filename)) {
    mf::LogInfo("evgen::ProtoDUNETriggeredBeam::FindFile") << "File exists. Opening " << filename;
    /*theFile = new TFile(filename.c_str());
    if (!theFile ||theFile->IsZombie() || !theFile->IsOpen()) {
      delete theFile;
      theFile = 0x0;
      throw cet::exception("ProtoDUNECalibration.cxx") << "Could not open " << filename;
    }*/
    return filename;
  }
  else {
    mf::LogInfo("evgen::ProtoDUNETriggeredBeam::FindFile") << "File does not exist here. Searching FW_SEARCH_PATH";
    cet::search_path sp{"FW_SEARCH_PATH"};
    std::string found_filename;
    auto found = sp.find_file(filename, found_filename);
    if (!found) {
      throw cet::exception("ProtoDUNECalibration.cxx") << "Could not find " << filename;
    }

    mf::LogInfo("evgen::ProtoDUNETriggeredBeam::FindFile") << "Found file " << found_filename;
    /*
    theFile = new TFile(found_filename.c_str());
    if (!theFile ||theFile->IsZombie() || !theFile->IsOpen()) {
      delete theFile;
      theFile = 0x0;
      throw cet::exception("ProtoDUNECalibration.cxx") << "Could not open " << found_filename;
    }*/
    return found_filename;
  }
}


DEFINE_ART_MODULE(evgen::ProtoDUNETriggeredBeam)