#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TBits.h>

#include <IHandle.hxx>
#include <IG4Trajectory.hxx>

#include <cometEventLoop.hxx>

const int kNPmax = 1000;
static TObjString monitor_name("Test");

class IExtractFromPlane: public COMET::ICOMETEventLoopFunction {
public:
    IExtractFromPlane() {
        fLSOption = "";
        fQuiet = false;
        fPrintList = false;
  
  fNormal.SetXYZ(0,0,1);
  fMonitorPos.SetXYZ(0,0,0);
  fMonitorRadius = 70;

  // Initialise some of the branches
  brEvtFlags = 0;
  brEvtCode = 0;
  brMonitorName = new TObjArray();     // the name of the monitor
  
  fOutputFile = new TFile("rootracker_output.root", "RECREATE");
  fTree = new TTree("gRooTracker", "SimG4 event tree rootracker format");

  fTree->Branch("EvtFlags", "TBits",      &brEvtFlags, 32000, 1);           
  fTree->Branch("EvtCode",  "TObjString", &brEvtCode,  32000, 1);           
  fTree->Branch("EvtNum",          &brEvtNum,          "EvtNum/I");             
  fTree->Branch("EvtXSec",         &brEvtXSec,         "EvtXSec/D");           
  fTree->Branch("EvtDXSec",        &brEvtDXSec,        "EvtDXSec/D");           
  fTree->Branch("EvtWght",         &brEvtWght,         "EvtWght/D");           
  fTree->Branch("EvtProb",         &brEvtProb,         "EvtProb/D");           
  fTree->Branch("EvtVtx",           brEvtVtx,          "EvtVtx[4]/D");             
  fTree->Branch("StdHepN",         &brStdHepN,         "StdHepN/I");             
  fTree->Branch("StdHepPdg",        brStdHepPdg,       "StdHepPdg[StdHepN]/I"); 
  fTree->Branch("StdHepStatus",     brStdHepStatus,    "StdHepStatus[StdHepN]/I");
  fTree->Branch("StdHepRescat",     brStdHepRescat,    "StdHepRescat[StdHepN]/I");
  fTree->Branch("StdHepX4",         brStdHepX4,        "StdHepX4[StdHepN][4]/D");
  fTree->Branch("StdHepP4",         brStdHepP4,        "StdHepP4[StdHepN][4]/D");
  fTree->Branch("StdHepPolz",       brStdHepPolz,      "StdHepPolz[StdHepN][3]/D");
  fTree->Branch("StdHepFd",         brStdHepFd,        "StdHepFd[StdHepN]/I");
  fTree->Branch("StdHepLd",         brStdHepLd,        "StdHepLd[StdHepN]/I");
  fTree->Branch("StdHepFm",         brStdHepFm,        "StdHepFm[StdHepN]/I");
  fTree->Branch("StdHepLm",         brStdHepLm,        "StdHepLm[StdHepN]/I");
  
  // Added for COMET
  fTree->Branch("Weight",         brWeight,        "Weight[StdHepN]/D");
  fTree->Branch("MonitorName", "TObjArray", brMonitorName);
    }

  virtual ~IExtractFromPlane() {};

    void Usage(void) {
        std::cout << "    -O quiet       Don't call TObject::ls()"
                  << std::endl;
        std::cout << "    -O list        List events in the file"
                  << std::endl;
  std::cout << "    -O xpos=[val]  x-position of the monitor in global coordinates [mm] (default: 0)"
            << std::endl;
  std::cout << "    -O ypos=[val]  y-position of the monitor in global coordinates [mm] (default: 0)"
            << std::endl;
  std::cout << "    -O zpos=[val]  z-position of the monitor in global coordinates [mm] (default: 0)"
            << std::endl;
  std::cout << "    -O xnorm=[val]  x value of the normal vector of the monitor in global coordinates (default: 0)"
            << std::endl;
  std::cout << "    -O ynorm=[val]  y value of the normal vector of the monitor in global coordinates (default: 0)"
            << std::endl;
  std::cout << "    -O znorm=[val]  z value of the normal vector of the monitor in global coordinates (default: 1)"
            << std::endl;
  std::cout << "    -O radius=[val] radius of the monitor [mm] (default: 70 mm)"
            << std::endl;
        std::cout << "    -O <option>    Set the option for TObject::ls()"
                  << std::endl;
    }

    virtual bool SetOption(std::string option,std::string value="") {
      //        if (value != "") return false;
        if (option == "quiet") fQuiet = true;
        else if (option == "list") {
            fPrintList = true;
        }
  else if (option == "xpos") {
    COMETVerbose("Setting x position of monitor to " << value);
    fMonitorPos.SetX(atof(value.c_str()));
  }
  else if (option == "ypos") {
    COMETVerbose("Setting y position of monitor to " << value);
    fMonitorPos.SetY(atof(value.c_str()));
  }
  else if (option == "zpos") {
    COMETVerbose("Setting z position of monitor to " << value);
    fMonitorPos.SetZ(atof(value.c_str()));
  }
  else if (option == "xnorm") {
    COMETVerbose("Setting x value of the normal vector of monitor to " << value);
    fNormal.SetX(atof(value.c_str()));
  }
  else if (option == "ynorm") {
    COMETVerbose("Setting y value of the normal vector of monitor to " << value);
    fNormal.SetY(atof(value.c_str()));
  }
  else if (option == "znorm") {
    COMETVerbose("Setting z value of the normal vector of monitor to " << value);
    fNormal.SetZ(atof(value.c_str()));
  }
  else if (option == "radius") {
    COMETVerbose("Setting radius of monitor to " << value);
    fMonitorRadius = atof(value.c_str());
  }

        else fLSOption = option;
        return true;
    }

    bool operator () (COMET::ICOMETEvent& event) {
        // reset the particle counter
        fParticleCounter = 0;

  std::cout << "Test: Event ID = " << event.GetEventId() << std::endl; //event.ls(fLSOption.c_str());

  COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories = event.Get<COMET::IG4TrajectoryContainer>("truth/G4Trajectories");
  if (trajectories) {
    std::cout<< trajectories->GetName() << std::endl;

    // Loop through all the trajectories in this container
    for (COMET::IG4TrajectoryContainer::const_iterator iter = trajectories->begin(); // an <int, IG4Trajectory> pair
         iter != trajectories->end(); ++iter) {
      
      const COMET::IG4Trajectory* aTraj = &(iter->second);
      std::cout << "Particle #" << iter->first << ": "; aTraj->ls();
      bool passed = PassesPlane(aTraj);
      std::cout << std::endl;
    }
  }
  else {
    std::cout << "Element is not found" << std::endl;
  }

  // Can now save this entry to the tree
  brEvtNum = event.GetEventId() + 1; // store the event number (starting at 1)
  brEvtXSec = -999;
  brEvtDXSec = -999;
  brEvtWght = -999;
  brEvtProb = -999;
  brEvtVtx[0] = 0; brEvtVtx[1] = 0; brEvtVtx[2] = 0; brEvtVtx[3] = 0; // initialise to zero since I think this is the origin of the event

  brStdHepN = fParticleCounter;
  fTree->Fill();
  fTree->AutoSave();

  brMonitorName->Clear();

        if (fPrintList) {
            COMETLog("Event " << event.GetContext());
        }
        return true;
    }

  // PassesPlane(const COMET::IG4Trajectory* theTraj)
  // -- Checks to see if the line between two points intersects a plane
  // -- Returns: true if it intersects and false if it doesn't
  bool PassesPlane(const COMET::IG4Trajectory* theTraj) {
      const std::vector<COMET::IG4TrajectoryPoint> theTrajPoints = theTraj->GetTrajectoryPoints();
      const TLorentzVector start = (theTrajPoints.begin())->GetPosition();
      const TLorentzVector end = (theTrajPoints.end()-1)->GetPosition();
      TVector3 start_pos(start.X(), start.Y(), start.Z());
      TVector3 end_pos(end.X(), end.Y(), end.Z());

      fNormal = fNormal.Unit(); // normalise the normal vector

      COMETVerbose("IExtractFromPlane::PassesPlane()");
      COMETVerbose("================================");
      COMETVerbose("Start Pos: (" << start_pos.x() << ", " << start_pos.y() << ", " << start_pos.z() << ") mm");
      COMETVerbose("End Pos: (" << end_pos.x() << ", " << end_pos.y() << ", " << end_pos.z() << ") mm");
      COMETVerbose("Monitor: ");
      COMETVerbose("\tNormal: (" << fNormal.x() << ", " << fNormal.y() << ", " << fNormal.z() << ")");
      COMETVerbose("\tPosition: (" << fMonitorPos.x() << ", " << fMonitorPos.y() << ", " << fMonitorPos.z() << ") mm");
      COMETVerbose("\tRadius: " << fMonitorRadius << " mm\n");

      // We need two vectors:
      TVector3 path = end_pos - start_pos; // the vector between the start and end positions
      TVector3 monitor_pos_rel_start = fMonitorPos - start_pos; // and the vector for the position of the monitor, relative to the start position
      COMETVerbose("Path between start and end: (" << path.x() << ", " << path.y() << ", " << path.z() << ")");
      COMETVerbose("Monitor Position (rel. to start pos): (" << monitor_pos_rel_start.x() << ", " << monitor_pos_rel_start.y() << ", " << monitor_pos_rel_start.z() << ")");
      
      // Calculate the scalar products of these with the normal
      // If they are 0 then they are perpendicular to the normal (i.e. parallel to the monitor)
      double scalar_prod_path = path.Dot(fNormal);
      double scalar_prod_start_pos = fNormal.Dot(monitor_pos_rel_start);
      COMETVerbose("Scalar Products between:");
      COMETVerbose("\tpath and normal: " << scalar_prod_path);
      COMETVerbose("\tstart pos and normal: " << scalar_prod_start_pos << "\n");


      // If both scalar products are zero then the path is in the monitor
      if (scalar_prod_path == 0 && scalar_prod_start_pos == 0) {
  COMETLog("Path is in the monitor");
  return false; 
      }
      // else if only scalar_prod_path is zero, then the path is parallel to the monitor but not in it
      else if (scalar_prod_path == 0 && scalar_prod_start_pos != 0) {
  COMETVerbose("Path is parallel to monitor and will never intersect");
  return false;
      }
      // else the monitor does intersect at one point on the path
      else {
  double steps = scalar_prod_start_pos / scalar_prod_path; // should be between 0 and 1
  COMETVerbose("Steps from start position to intersection: " << steps);

  // Need to ensure that the number of steps is between 0 and 1 otherwise the intersection point would be beyond/behind the path
  if (steps < 0 || steps > 1) {
    COMETVerbose("Intersection point with plane is beyond/behind the path taken");
    return false;
  }
  // if the particle ends in the monitor
  if (steps == 1) {
    // ignore it so that we avoid double counting
    COMETVerbose("Particle ended in the monitor. Ignoring so that a double count is avoided");
    return false;
  }


  // Get the intersection point relative to the monitor position
  TVector3 intersection_point = steps*path + start_pos;
  TVector3 intersection_point_rel_monitor = intersection_point - fMonitorPos;

  // Check that the intersection point is within the physical monitor
  double r_intersect = intersection_point_rel_monitor.Perp();
  if (r_intersect <= fMonitorRadius) {
    COMETVerbose("Intersection!");
    COMETVerbose("Intersection point (global): (" << intersection_point.x() << ", " << intersection_point.y() << ", " << intersection_point.z() << ")");
    COMETVerbose("Intersection point (local): (" << intersection_point_rel_monitor.x() << ", " << intersection_point_rel_monitor.y() << ", " << intersection_point_rel_monitor.z() << ")");
    OutputRooTrackerEvent(theTraj, intersection_point_rel_monitor);
    ++fParticleCounter;
    return true;
  }
  else {
    COMETVerbose("No Intersection!");
    return false;
  }
      }
    }

  void OutputRooTrackerEvent(const COMET::IG4Trajectory* theTraj, TVector3 intersection_point_rel_monitor) {

    // NB Need to convert:
    // from mm --> m
    // from ?eV --> MeV
      brStdHepPdg[fParticleCounter] = theTraj->GetPDGEncoding();
      brStdHepStatus[fParticleCounter] = 1; // set to 1 since that is what SimG4 is expecting for particles that it can track
      brStdHepRescat[fParticleCounter] = -999;  
      brStdHepX4    [fParticleCounter][0] = intersection_point_rel_monitor.x() / 1000;  
      brStdHepX4    [fParticleCounter][1] = intersection_point_rel_monitor.y() / 1000; 
      brStdHepX4    [fParticleCounter][2] = intersection_point_rel_monitor.z() / 1000; 
      brStdHepX4    [fParticleCounter][3] = 0;
      //      brStdHepX4    [fParticleCounter][3] = particle_info_.ctoffc;
      
      brStdHepP4    [fParticleCounter][0] = 0;
      brStdHepP4    [fParticleCounter][1] = 0;
      brStdHepP4    [fParticleCounter][2] = 0;
      brStdHepP4    [fParticleCounter][3] = 0;
      //      brStdHepP4    [fParticleCounter][3] = particle_info_.etc;

      brStdHepPolz  [fParticleCounter][0] = -999;
      brStdHepPolz  [fParticleCounter][1] = -999;
      brStdHepPolz  [fParticleCounter][2] = -999;
      
      brStdHepFd    [fParticleCounter] = -999;     // First daughter
      brStdHepLd    [fParticleCounter] = -999;     // Last  daughter
      brStdHepFm    [fParticleCounter] = theTraj->GetParentId();     // First mother
      brStdHepLm    [fParticleCounter] = -999;     // Last  mother
      
      brWeight      [fParticleCounter] = 1;	

      brMonitorName->Add(&monitor_name);
    }

    // Do nothing... This is here to test compiler warnings.  The warning
    // can be prevented by adding
    //
    // using COMET::ICOMETEventLoopFunction::Finalize;
    void Finalize(COMET::ICOMETOutput*const output) {
    }

private:
    std::string fLSOption;
    bool fQuiet;
    bool fPrintList;

    TVector3 fNormal; // assuming normal points to the "front" of the monitor
    TVector3 fMonitorPos;
    double fMonitorRadius;

    TFile* fOutputFile;
    TTree* fTree;
  
  // RooTracker Branches
  TBits*      brEvtFlags; // Generator-specific event flags
  TObjString* brEvtCode; // Generator-specific event string with 'event code'
  int         brEvtNum;       // Event number
  double      brEvtXSec;      // Cross section for selected event (1E-38 cm2)
  double      brEvtDXSec;     // Cross section for selected event (1E-38 cm2 /{K^n})
  double      brEvtWght;      // Weight of the event
  double      brEvtProb;      // Probability for that event (gicen cross-sections, path lengths etc.)
  double      brEvtVtx[4];    // Event vertex position in detector coordinate system (SI)
  int         brStdHepN;      // Number of particles in particle array
  // std-hep like particle array
  int         brStdHepPdg   [kNPmax];     // Pdg codes (& generator specific codes for pseudoparticles)
  int         brStdHepStatus[kNPmax];     // Generator-specific status code
  int         brStdHepRescat[kNPmax];     // Hadron transport model - specific rescattering code
  double      brStdHepX4    [kNPmax][4];  // 4-x (x, y, z, t) of particle in hit nucleus frame (fm)
  double      brStdHepP4    [kNPmax][4];  // 4-p (px,py,pz,E) of particle in LAB frame (GeV)
  double      brStdHepPolz  [kNPmax][3];  // Polarization vector
  int         brStdHepFd    [kNPmax];     // First daughter
  int         brStdHepLd    [kNPmax];     // Last  daughter
  int         brStdHepFm    [kNPmax];     // First mother
  int         brStdHepLm    [kNPmax];     // Last  mother

  // Added for COMET
  double      brWeight      [kNPmax];     // weight of the particle
  TObjArray*  brMonitorName;     // the name of the monitor

  int fParticleCounter; // keep track of the number of particles that pass through the monitor
};

int main(int argc, char **argv) {
    IExtractFromPlane userCode;
    COMET::cometEventLoop(argc,argv,userCode,1);
}