#include <iostream>
#include <cmath>
#include <algorithm>
using namespace std;

#include <TLorentzVector.h>
#include <TMath.h>

#include <JANA/JApplication.h>
#include <JANA/JEventLoop.h>
using namespace jana;

#include <TRACKING/DMCThrown.h>
#include <TRACKING/DTrackTimeBased.h>
#include <FCAL/DFCALShower.h>
#include <BCAL/DBCALShower.h>
#include <units.h>

#include "DEventProcessor_photon_discrim.h"
//==========================================================================
// This file contains code for a plugin that will create a few histograms
// based on the reconstructed data. It can mainly serve as an example
// of how one could access the reconstructed data in their own analysis
// program.
//
// Histograms are created in the init() method and they are filled in
// the evnt() method.
//==========================================================================

// Routine used to create our JEventProcessor
extern "C"{
  void InitPlugin(JApplication *app){
    InitJANAPlugin(app);
    app->AddProcessor(new DEventProcessor_photon_discrim());
  }
}

//------------------
// init
//------------------
jerror_t DEventProcessor_photon_discrim::init(void)
{

  photon_discrim_BCAL = new TTree("photon_discrim_BCAL","photon_discrim_BCAL");

  photon_discrim_BCAL->Branch("recon_BCAL_match_thrown",&m_recon_BCAL_match_thrown,"recon_BCAL_match_thrown/O");
  photon_discrim_BCAL->Branch("recon_BCAL_z_entry",&m_recon_BCAL_z_entry,"recon_BCAL_z_entry/F");
  photon_discrim_BCAL->Branch("recon_BCAL_z_cluster",&m_recon_BCAL_z_cluster,"recon_BCAL_z_cluster/F");
  photon_discrim_BCAL->Branch("recon_BCAL_r",&m_recon_BCAL_r,"recon_BCAL_r/F");
  photon_discrim_BCAL->Branch("recon_BCAL_phi",&m_recon_BCAL_phi,"recon_BCAL_phi/F");
  photon_discrim_BCAL->Branch("recon_BCAL_tproj",&m_recon_BCAL_tproj,"recon_BCAL_tproj/F");
  photon_discrim_BCAL->Branch("recon_BCAL_E",&m_recon_BCAL_E,"recon_BCAL_E/F");
  photon_discrim_BCAL->Branch("recon_BCAL_closest_shower_dist",&m_recon_BCAL_closest_shower_dist,"recon_BCAL_closest_shower_dist/F");
  photon_discrim_BCAL->Branch("recon_BCAL_closest_shower_E",&m_recon_BCAL_closest_shower_E,"recon_BCAL_closest_shower_E/F");
  photon_discrim_BCAL->Branch("recon_BCAL_closest_track_dist",&m_recon_BCAL_closest_track_dist,"recon_BCAL_track_shower_dist/F");

  photon_discrim_tree = new TTree("photon_discrim_tree","photon_discrim_tree");

  photon_discrim_tree->Branch("n_recon_FCAL",&m_n_recon_FCAL,"n_recon_FCAL/I");
  photon_discrim_tree->Branch("recon_FCAL_tproj",m_recon_FCAL_tproj,"recon_FCAL_tproj[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_x",m_recon_FCAL_x,"recon_FCAL_x[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_y",m_recon_FCAL_y,"recon_FCAL_y[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_z",m_recon_FCAL_z,"recon_FCAL_z[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_E",m_recon_FCAL_E,"recon_FCAL_E[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_RMS",m_recon_FCAL_RMS,"recon_FCAL_RMS[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_RMS_x",m_recon_FCAL_RMS_x,"recon_FCAL_RMS_x[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_RMS_y",m_recon_FCAL_RMS_y,"recon_FCAL_RMS_y[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_RMS_u",m_recon_FCAL_RMS_u,"recon_FCAL_RMS_u[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_RMS_v",m_recon_FCAL_RMS_v,"recon_FCAL_RMS_v[n_recon_FCAL]/F");
  photon_discrim_tree->Branch("recon_FCAL_splash",m_recon_FCAL_splash,"recon_FCAL_splash[n_recon_FCAL]/F");

  photon_discrim_tree->Branch("n_thrown_photons",&m_n_thrown_photons,"n_thrown_photons/I");
  photon_discrim_tree->Branch("thrown_E",m_thrown_E,"thrown_E[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_px",m_thrown_px,"thrown_px[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_py",m_thrown_py,"thrown_py[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_pz",m_thrown_pz,"thrown_pz[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_vx",m_thrown_vx,"thrown_vx[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_vy",m_thrown_vy,"thrown_vy[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_vz",m_thrown_vz,"thrown_vz[n_thrown_photons]/F");
  photon_discrim_tree->Branch("thrown_BCAL",m_thrown_BCAL,"thrown_BCAL[n_thrown_photons]/O");
  photon_discrim_tree->Branch("thrown_BCAL_z_entry",m_thrown_BCAL_z_entry,"thrown_BCAL_z_entry[n_thrown_photons]/F");

  pthread_mutex_init(&mutex, NULL);

  return NOERROR;
}

//------------------
// evnt 
//------------------
jerror_t DEventProcessor_photon_discrim::evnt(JEventLoop *loop, int eventnumber)
{
  vector<const DMCThrown*> throwns;
  vector<const DBCALShower*> bcal_showers;
  vector<const DFCALShower*> fcal_showers;
  //should I use DChargedTrackHypothesis instead
  vector<const DTrackTimeBased*> tracks;

  loop->Get(throwns);
  loop->Get(bcal_showers);
  loop->Get(fcal_showers);
  loop->Get(tracks);

  //--------------------------------------------------------------------
  // Fill histograms below here
  pthread_mutex_lock(&mutex);

  //reset tree stuff
  m_n_recon_FCAL=min((int)fcal_showers.size(),(int)MAX_PHOTON);
  m_n_thrown_photons=0; //we will count these later

  const double bcal_inner_r = 64.3; //should look this up

  //Get the thrown photons. Also the thrown vertex
  vector<const DMCThrown*> thrown_photons;
  for (unsigned int i=0; i<throwns.size();i++) {
    int type=throwns[i]->type;
    //need to check parentid to make sure photon isn't from a decay in geant
    //if the hddm file is from genr8, the condition for a primary particle
    //is: throwns[i]->parentid != 0
    //Otherwise, check for primary-ness by requiring that the vertex is close
    //to the beamline
    double vertex_r = sqrt(throwns[i]->x()*throwns[i]->x() + throwns[i]->y()*throwns[i]->y());
    if (type != 1 || vertex_r > .5) continue;
    //it's a photon!
    if (m_n_thrown_photons >= MAX_PHOTON) continue;
    thrown_photons.push_back(throwns[i]);

    m_thrown_E[m_n_thrown_photons] = throwns[i]->lorentzMomentum().E();
    m_thrown_px[m_n_thrown_photons] = throwns[i]->lorentzMomentum().X();
    m_thrown_py[m_n_thrown_photons] = throwns[i]->lorentzMomentum().Y();
    m_thrown_pz[m_n_thrown_photons] = throwns[i]->lorentzMomentum().Z();
    m_thrown_vx[m_n_thrown_photons] = throwns[i]->position().X();
    m_thrown_vy[m_n_thrown_photons] = throwns[i]->position().Y();
    m_thrown_vz[m_n_thrown_photons] = throwns[i]->position().Z();

    //We wish to calculate the z_entry of the thrown photon
    //(the point when the photon would hit the inner radius of the BCAL)
    //We have the vertex and the momentum of the thrown photon
    //so the path of the photon can be parameterized as vector(vertex) + s*vector(momentum), where s is positive. we want to find the value of s such that x^2+y^2=inner_rad^2, and then use that value of s to find the z where the photon enters the BCAL. the rest is just algebra...
    DVector3 vertex = throwns[i]->position();
    DLorentzVector momentum = throwns[i]->lorentzMomentum();
    double a = momentum.X()*momentum.X()+momentum.Y()*momentum.Y();
    double b = 2*vertex.X()*momentum.X()+2*vertex.Y()*momentum.Y();
    double c = vertex.X()*vertex.X()+vertex.Y()*vertex.Y()-bcal_inner_r*bcal_inner_r;
    double s = (-b+sqrt(b*b-4*a*c))/(2*a);
    double thrown_z_entry = s*momentum.Z()+vertex.Z();

    //should look these numbers up
    if (thrown_z_entry > 17. && thrown_z_entry < 407.) {
      m_thrown_BCAL[m_n_thrown_photons] = true;
      m_thrown_BCAL_z_entry[m_n_thrown_photons] = thrown_z_entry;
    } else {
      m_thrown_BCAL[m_n_thrown_photons] = false;
      m_thrown_BCAL_z_entry[m_n_thrown_photons] = 0.;
    }

    m_n_thrown_photons++;
  }

  for (int i=0; i<(int)bcal_showers.size(); i++) {
    m_recon_BCAL_match_thrown = false; //for now just assume no match

    const double target_center_z = 65.0;

    double x = bcal_showers[i]->x;
    double y = bcal_showers[i]->y;
    double z = bcal_showers[i]->z;
    double r = sqrt(x*x + y*y);
    double phi = atan2(y,x);
    m_recon_BCAL_z_cluster = bcal_showers[i]->z; //this is the z of the DBCALShower, not at a fixed radius
    m_recon_BCAL_r = r;
    m_recon_BCAL_phi = phi;

    //to calculate z_entry, just assume that the shower is caused by a shower
    //originating at (0,0,65)
    double z_entry = (z-target_center_z)*(bcal_inner_r/r)+target_center_z;
    m_recon_BCAL_z_entry = z_entry; //position where photon hits inner radius of BCAL

    //again assume we have a photon originating at (0,0,65) at t=0
    //this makes the assumption that we always identify the correct RF bunch
    m_recon_BCAL_tproj = bcal_showers[i]->t - sqrt( (z-target_center_z)*(z-target_center_z) + r*r)/SPEED_OF_LIGHT;
    m_recon_BCAL_E = bcal_showers[i]->E;

    //check if this matches any of the thrown photons
    for (int j=0; j<(int)thrown_photons.size(); j++) {
      //only consider thrown photons that should end up in the BCAL.
      if (!m_thrown_BCAL[j]) continue;
      double thrown_phi = thrown_photons[j]->lorentzMomentum().Phi();
      double recon_phi = phi;
      double phi_diff = thrown_phi - recon_phi;
      double phi_diff_alt = (thrown_phi > recon_phi ?
                             thrown_phi - recon_phi - 2*M_PI :
                             recon_phi - thrown_phi - 2*M_PI);
      phi_diff = min(fabs(phi_diff), fabs(phi_diff_alt));

      double thrown_E = thrown_photons[j]->lorentzMomentum().E();
      double recon_E = bcal_showers[i]->E;
      double delta_E = recon_E - thrown_E;
      double sig_E_NIM = thrown_E*sqrt(.054*.054/thrown_E + .023*.023);
      double delta_E_thresh = min(3.0*sig_E_NIM,thrown_E);

      //We have already calculated the z_entry of the thrown photon
      //(the point when the photon would hit the inner radius of the BCAL)
      double thrown_z_entry = m_thrown_BCAL_z_entry[j];
      //recalculate z_entry using thrown vertex
      //since we require the thrown vertex to be very close to beamline,
      //simplify the calculation by treating vertex_x, vertex_y as zero
      double recon_z_entry = (z-m_thrown_vz[j])*(bcal_inner_r/r)+m_thrown_vz[j];
      double delta_z = recon_z_entry - thrown_z_entry;
      double sig_z_NIM = 1.1/sqrt(thrown_E);
      double delta_z_thresh = min(5.0*sig_z_NIM,39.0);
      delta_z_thresh = max(delta_z_thresh,9.0);

      //should adjust these numbers
      if (fabs(delta_E) < delta_E_thresh
          && phi_diff < .07
          && fabs(delta_z) < delta_z_thresh) {
        m_recon_BCAL_match_thrown = true;
      }
    }

    //Find the closest BCAL shower to this one
    //This is an inefficient way to do this
    double min_sep = 500;
    double closest_E = 0;
    for (int j=0; j<(int)bcal_showers.size(); j++) {
      if (i==j) continue; //make sure we're not comparing the shower to itself

      double x2 = bcal_showers[j]->x;
      double y2 = bcal_showers[j]->y;
      double z2 = bcal_showers[j]->z;
      double dist = sqrt((x-x2)*(x-x2)+(y-y2)*(y-y2)+(z-z2)*(z-z2)); //for a first past just use a Cartesian distance
      if (dist<min_sep) {
        min_sep = dist;
        closest_E = bcal_showers[j]->E;
      }
    }

    m_recon_BCAL_closest_shower_dist = min_sep;
    m_recon_BCAL_closest_shower_E = closest_E;

    //Find closest track
    //This is inefficient
    double min_sep_track = 500;
    for (int j=0; j<(int)tracks.size(); j++) {
      const DReferenceTrajectory* rt_recon = tracks[j]->rt;
      double path_length=0,flight_time=0;
      TVector3 bcal_pos(bcal_showers[i]->x,bcal_showers[i]->y,bcal_showers[i]->z);
      double d = rt_recon->DistToRTwithTime(bcal_pos,&path_length,&flight_time);
      //could also consider looking at flight time


      if (d < min_sep_track) {
        min_sep_track = d;
      }
    }
    m_recon_BCAL_closest_track_dist = min_sep_track;

    photon_discrim_BCAL->Fill();

  }

  for (int i=0; i<m_n_recon_FCAL; i++) {
    m_recon_FCAL_match_thrown[i] = false;
    m_recon_FCAL_tproj[i] = fcal_showers[i]->getTime();
    m_recon_FCAL_x[i] = fcal_showers[i]->getPosition().X();
    m_recon_FCAL_y[i] = fcal_showers[i]->getPosition().Y();
    m_recon_FCAL_z[i] = fcal_showers[i]->getPosition().Z();
    m_recon_FCAL_E[i] = fcal_showers[i]->getEnergy();
    m_recon_FCAL_RMS[i] = 0;
    m_recon_FCAL_RMS_x[i] = 0;
    m_recon_FCAL_RMS_y[i] = 0;
    m_recon_FCAL_RMS_u[i] = 0;
    m_recon_FCAL_RMS_v[i] = 0;
    m_recon_FCAL_splash[i] = 0;
  }

  /*const double phiDiffFCAL=.07;
  const double thetaDiffFCAL=.5/radToDeg;
  const double phiDiffBCAL=.07; //the width of one BCAL sector (inner layer) is about .03 rad

  //only match one reconstructed photon with each thrown photon
  bool matched[NTHROWNPHOTONS]={false,false};
  vector<DLorentzVector> photon_OK; //photons that pass some cut
  for (unsigned int i=0; i < neutrals.size(); i++) {
    //only keep neutrals if they are photons
    if (neutrals[i]->PID() != Gamma ) continue;

    //which detector detected this particle?
    DetectorSystem_t calorimeter=SYS_NULL;
    vector<const DNeutralShower*> candidates;
    neutrals[i]->Get(candidates);
    const DFCALCluster *fcal_cluster=NULL;
    const DBCALShower *bcal_shower=NULL;
    if (candidates.size()==1) {
      calorimeter = candidates[0]->dDetectorSystem;
      //while we're at it, also get the associated DFCALCluster or DBCALShower, which will be used below
      if (calorimeter == SYS_FCAL) {
        vector<const DFCALShower*> fcal_shower;
        candidates[0]->Get(fcal_shower);
        if (fcal_shower.size()==1) {
          vector<const DFCALCluster*> fcal_cluster_vect;
          fcal_shower[0]->Get(fcal_cluster_vect);
          if (fcal_cluster_vect.size()==1) fcal_cluster=fcal_cluster_vect[0];
        }
      } else if (calorimeter == SYS_BCAL) {
        vector<const DBCALShower*> bcal_shower_vect;
        candidates[0]->Get(bcal_shower_vect);
        if (bcal_shower_vect.size()==1) bcal_shower=bcal_shower_vect[0];
      }
    }

    const DNeutralParticleHypothesis *photon=neutrals[i];

    //calculate z for BCAL photons. this is the z where the photon hits the inner radius, assuming the photon goes into the BCAL.

    //so the path of the photon can be parameterized as vector(vertex) + s*vector(momentum), where s is positive. we want to find the value of s such that x^2+y^2=inner_rad^2, and then use that value of s to find the z where the photon enters the BCAL. the rest is just algebra...
    DVector3 vertex = photon->position();
    DLorentzVector momentum = photon->lorentzMomentum();
    double a = momentum.X()*momentum.X()+momentum.Y()*momentum.Y();
    double b = 2*vertex.X()*momentum.X()+2*vertex.Y()*momentum.Y();
    double c = vertex.X()*vertex.X()+vertex.Y()*vertex.Y()-inner_rad*inner_rad;
    double s = (-b+sqrt(b*b-4*a*c))/(2*a);
    double z = s*momentum.Z()+vertex.Z();
    //these z errors are determined by fitting with single particles. svn8513. 1234 summing.
    //I think this is not quite right. Check whether the z_error I calculate uses raw energies or corrected ones?
    //double z_err = 2.34/sqrt(photon->lorentzMomentum().E()); //KLOE
    double z_err = 1.17/sqrt(photon->lorentzMomentum().E())+.17/photon->lorentzMomentum().E(); //svn8513+newzerr

    photon_OK.push_back(photon->lorentzMomentum());

    if (calorimeter==SYS_FCAL) {
      photonEvsTheta_recon_FCAL->Fill(photon->lorentzMomentum().Theta()*radToDeg,photon->lorentzMomentum().E());

      if (m_n_recon_FCAL < MAX_PHOTON) {
        m_recon_FCAL_tproj[m_n_recon_FCAL] = photon->TOF() - photon->pathLength()/SPEED_OF_LIGHT;
        m_recon_FCAL_px[m_n_recon_FCAL]=photon->lorentzMomentum().X();
        m_recon_FCAL_py[m_n_recon_FCAL]=photon->lorentzMomentum().Y();
        m_recon_FCAL_pz[m_n_recon_FCAL]=photon->lorentzMomentum().Z();
        m_recon_FCAL_E[m_n_recon_FCAL]=photon->lorentzMomentum().E();
        //cluster shape
        if (fcal_cluster!=NULL) {
          m_recon_FCAL_RMS[m_n_recon_FCAL]=fcal_cluster->getRMS();
          m_recon_FCAL_RMS_x[m_n_recon_FCAL]=fcal_cluster->getRMS_x();
          m_recon_FCAL_RMS_y[m_n_recon_FCAL]=fcal_cluster->getRMS_y();
          m_recon_FCAL_RMS_u[m_n_recon_FCAL]=fcal_cluster->getRMS_u();
          m_recon_FCAL_RMS_v[m_n_recon_FCAL]=fcal_cluster->getRMS_v();
          m_recon_FCAL_splash[m_n_recon_FCAL]=fcal_cluster->splash;
        } else {
          m_recon_FCAL_RMS[m_n_recon_FCAL]=0;
          m_recon_FCAL_RMS_x[m_n_recon_FCAL]=0;
          m_recon_FCAL_RMS_y[m_n_recon_FCAL]=0;
          m_recon_FCAL_RMS_u[m_n_recon_FCAL]=0;
          m_recon_FCAL_RMS_v[m_n_recon_FCAL]=0;
          m_recon_FCAL_splash[m_n_recon_FCAL]=0;
        }
        m_n_recon_FCAL++;
      }

    } else if (calorimeter==SYS_BCAL) {
      

      photonEvsTheta_recon_BCAL->Fill(photon->lorentzMomentum().Theta()*radToDeg,photon->lorentzMomentum().E());
      photonEvsZ_recon_BCAL->Fill(z,photon->lorentzMomentum().E());

      if (m_n_recon_BCAL < MAX_PHOTON) {
        m_recon_BCAL_z[m_n_recon_BCAL]=z;
        m_recon_BCAL_tproj[m_n_recon_BCAL] = photon->TOF() - photon->pathLength()/SPEED_OF_LIGHT;
        m_recon_BCAL_px[m_n_recon_BCAL]=photon->lorentzMomentum().X();
        m_recon_BCAL_py[m_n_recon_BCAL]=photon->lorentzMomentum().Y();
        m_recon_BCAL_pz[m_n_recon_BCAL]=photon->lorentzMomentum().Z();
        m_recon_BCAL_E[m_n_recon_BCAL]=photon->lorentzMomentum().E();
        if (bcal_shower!=NULL) {
          double r = sqrt(bcal_shower->x*bcal_shower->x+bcal_shower->y*bcal_shower->y);
          m_recon_BCAL_z_cluster[m_n_recon_BCAL]=bcal_shower->z;
          m_recon_BCAL_r[m_n_recon_BCAL]=r;
        } else {
          m_recon_BCAL_z_cluster[m_n_recon_BCAL]=0;
          m_recon_BCAL_r[m_n_recon_BCAL]=0;
        }
        m_n_recon_BCAL++;
      }
    }

    for (unsigned int j=0; j<thrown_photons.size();j++) {
      //only allow one matched photon for each thrown photon
      if (matched[j]) continue;

      double z_thrown = inner_rad/tan(thrown_photons[j].Theta())+thrownVertex.Z();

      //Does a photon in the FCAL match a thrown photon?
      //This depends on having a good vertex...
      if (fabs(thrown_photons[j].Theta()-photon->lorentzMomentum().Theta())<thetaDiffFCAL
          && fabs(thrown_photons[j].Phi()-photon->lorentzMomentum().Phi())<phiDiffFCAL
          && calorimeter==SYS_FCAL ) {

        matched[j]=true;
	  
        double thrownE = thrown_photons[j].E();
        double reconE = photon->lorentzMomentum().E();
        double errorE = sqrt(photon->errorMatrix()(3,3));

        gamma_E_res_FCAL->Fill((reconE-thrownE)/thrownE);
        gamma_E_res_norm_FCAL->Fill((reconE-thrownE)/errorE);

        photonEvsTheta_good_FCAL->Fill(photon->lorentzMomentum().Theta()*radToDeg,photon->lorentzMomentum().E());
      } else if ( fabs(z_thrown-z) < z_err*3.
                  && fabs(thrown_photons[j].Phi()-photon->lorentzMomentum().Phi()) < phiDiffBCAL
                  && calorimeter==SYS_BCAL ) {

        matched[j]=true;

        double thrownE = thrown_photons[j].E();
        double reconE = photon->lorentzMomentum().E();
        double errorE = sqrt(photon->errorMatrix()(3,3));

        gamma_E_res_BCAL->Fill((reconE-thrownE)/thrownE);
        gamma_E_res_BCAL_vs_E->Fill(thrownE,(reconE-thrownE)/thrownE);
        gamma_E_res_norm_BCAL->Fill((reconE-thrownE)/errorE);

        photonEvsTheta_good_BCAL->Fill(photon->lorentzMomentum().Theta()*radToDeg,photon->lorentzMomentum().E());
      }
    }
  }*/

  photon_discrim_tree->Fill();

  pthread_mutex_unlock(&mutex);

  return NOERROR;
}

//------------------
// erun
//------------------
jerror_t DEventProcessor_photon_discrim::erun(void)
{
  return NOERROR;
}

//------------------
// fini
//------------------
jerror_t DEventProcessor_photon_discrim::fini(void)
{
  // Histograms are automatically written to file by hd_root program (or equivalent)
  // so we don't need to do anything here.
  
  return NOERROR;
}