#include "DEventProcessor_bcal_charged_timing.h"

#include "DANA/DApplication.h"
#include "BCAL/DBCALShower.h"
#include "BCAL/DBCALCluster.h"
#include "BCAL/DBCALPoint.h"
#include "TRACKING/DMCThrown.h"
#include "PID/DChargedTrackHypothesis.h"
#include "PID/DNeutralParticleHypothesis.h"
#include "units.h"

#define DPRINT(x) cout << #x "=" << x << endl

// The executable should define the ROOTfile global variable. It will
// be automatically linked when dlopen is called.
extern TFile *ROOTfile;

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

bool TimeSort(const DBCALPoint* point1, const DBCALPoint* point2) {
  return point1->t() < point2->t();
}

//global variables used for sort function
//This will not work in multi-threaded running!!
//These coordinates are the "entry point" of the cluster. That is, if you draw
//a straight line from the center of the target to the cluster location, these
//are the coordinates of where that line intersects the inner radius of the BCAL.
double gl_x;
double gl_y;
double gl_z;

//calculate distance to entry point
double DistToEntry(const DBCALPoint *point1) {
  double x1 = point1->r()*cos(point1->phi());
  double y1 = point1->r()*sin(point1->phi());
  double z1 = point1->z();
  double dist_sq = (x1 - gl_x)*(x1 - gl_x) + (y1 - gl_y)*(y1 - gl_y) + (z1 - gl_z)*(z1 - gl_z);
  return sqrt(dist_sq);
}

bool PointToPointDistSort(const DBCALPoint* point1, const DBCALPoint* point2) {
  return DistToEntry(point1) < DistToEntry(point2);
}

//calculate shortest distance from point1 to line from center of target to cluster position
double DistToTrack(const DBCALPoint *point1) {
  //we want to find the distance between a point, and a line defined 
  //by the origin (the target center) (0,0,0) and the point (gl_x,gl_y,gl_z)
  //following http://mathworld.wolfram.com/Point-LineDistance3-Dimensional.html

  TVector3 x1(0.0,0.0,0.0);
  TVector3 x2(gl_x,gl_y,gl_z);

  double x1p= point1->r()*cos(point1->phi());
  double y1p = point1->r()*sin(point1->phi());
  double z1p = point1->z();
  TVector3 x0(x1p,y1p,z1p);
  return ((x0-x1).Cross(x0-x2)).Mag() / (x2-x1).Mag();
}

bool PointToLineDistSort(const DBCALPoint* point1, const DBCALPoint* point2) {
  return DistToTrack(point1) < DistToTrack(point2);
}

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

  //set up the tree
  bcal_charged_timing_tree = new TTree("bcal_charged_timing_tree","bcal_charged_timing_tree");

  bcal_charged_timing_tree->Branch("track_recon",&m_track_recon,"track_recon/O");
  bcal_charged_timing_tree->Branch("bcal_match",&m_bcal_match,"bcal_match/O");
  bcal_charged_timing_tree->Branch("cluster_t",&m_cluster_t,"cluster_t/F");
  bcal_charged_timing_tree->Branch("cluster_t_first",&m_cluster_t_first,"cluster_t_first/F");
  bcal_charged_timing_tree->Branch("cluster_t_early_half",&m_cluster_t_early_half,"cluster_t_early_half/F");
  bcal_charged_timing_tree->Branch("cluster_t_close_to_entry_half",&m_cluster_t_close_to_entry_half,"cluster_t_close_to_entry_half/F");
  bcal_charged_timing_tree->Branch("cluster_t_close_to_track_half",&m_cluster_t_close_to_track_half,"cluster_t_close_to_track_half/F");
  bcal_charged_timing_tree->Branch("cluster_t_grabbag",&m_cluster_t_grabbag,"cluster_t_grabbag/F");
  bcal_charged_timing_tree->Branch("cluster_E",&m_cluster_E,"cluster_E/F");
  bcal_charged_timing_tree->Branch("cluster_num_points",&m_cluster_num_points,"cluster_num_points/I");
  bcal_charged_timing_tree->Branch("track_t",&m_track_t,"track_t/F");

  associated_points_tree = new TTree("associated_points_tree","associated_points_tree");
  associated_points_tree->Branch("point_E",&m_point_E,"point_E/F");
  associated_points_tree->Branch("point_t",&m_point_t,"point_t/F");
  associated_points_tree->Branch("point_t_inner",&m_point_t_inner,"point_t_inner/F");
  associated_points_tree->Branch("point_z",&m_point_z,"point_z/F");
  associated_points_tree->Branch("point_dist_to_entry",&m_point_dist_to_entry,"point_dist_to_entry/F");
  associated_points_tree->Branch("point_dist_to_track",&m_point_dist_to_track,"point_dist_to_track/F");
  associated_points_tree->Branch("point_layer",&m_point_layer,"point_layer/I");

  return NOERROR;
}

//------------------
// brun
//------------------
jerror_t DEventProcessor_bcal_charged_timing::brun(JEventLoop *eventLoop, int runnumber)
{

  return NOERROR;
}

//------------------
// evnt
//------------------
jerror_t DEventProcessor_bcal_charged_timing::evnt(JEventLoop *loop, int eventnumber)
{
  vector<const DMCThrown*> mcthrowns;
  vector<const DChargedTrackHypothesis*> charged_tracks;
  vector<const DNeutralParticleHypothesis*> neutrals;
  loop->Get(mcthrowns);
  loop->Get(charged_tracks);
  loop->Get(neutrals);

  //the assumption here is that we are analyzing single-particle events
  //however there may be >1 mcthrowns due to in-flight decays or something
  if (mcthrowns.size() == 0) return NOERROR;
  //try to get this working for both photons and charged pions
  bool is_charged;
  if (mcthrowns[0]->PID() == PiPlus || mcthrowns[0]->PID() == PiMinus || mcthrowns[0]->PID() == Proton) {
    is_charged = true;
  } else if (mcthrowns[0]->PID() == Gamma) {
    is_charged = false;
  } else {
    return NOERROR;
  }
  DVector3 thrown_pos = mcthrowns[0]->position();
  DVector3 thrown_mom = mcthrowns[0]->momentum();

  //check if there is a "good_track" i.e. a reconstructed track (or photon) that matches the thrown PID and also the momentum, within some tolerance
  //this assumes that there is only one "good" track, which I hope is the case, at least most of the time
  const DKinematicData* good_track = NULL;
  if (is_charged) {
    for (int i=0; i<(int)charged_tracks.size(); i++) {
      if (charged_tracks[i]->PID() != mcthrowns[0]->PID()) continue;

      DVector3 recon_mom = charged_tracks[i]->momentum();
      //I should check if this matching criteria is optimal
      if (fabs(recon_mom.Mag() - thrown_mom.Mag())/thrown_mom.Mag() > .2) continue;
      if (fabs(recon_mom.Theta() - thrown_mom.Theta()) > .09) continue;
      if (fabs(recon_mom.Phi() - thrown_mom.Phi()) > .09) continue;

      good_track = charged_tracks[i];
    }
  } else {
    for (int i=0; i<(int)neutrals.size(); i++) {
      if (neutrals[i]->PID() != mcthrowns[0]->PID()) continue;

      DVector3 recon_mom = neutrals[i]->momentum();
      //I should check if this matching criteria is optimal
      if (fabs(recon_mom.Mag() - thrown_mom.Mag())/thrown_mom.Mag() > .2) continue;
      if (fabs(recon_mom.Theta() - thrown_mom.Theta()) > .09) continue;
      if (fabs(recon_mom.Phi() - thrown_mom.Phi()) > .09) continue;

      good_track = neutrals[i];
    }
  }

  LockState();

  //reset all these values from the last event
  m_track_recon = false;
  m_bcal_match = false;
  m_cluster_t = 0.0;
  m_cluster_E = 0.0;
  m_cluster_num_points = 0;
  m_cluster_t_first = 0.0;
  m_cluster_t_early_half = 0.0;
  m_cluster_t_close_to_entry_half = 0.0;
  m_cluster_t_close_to_track_half = 0.0;
  m_cluster_t_grabbag = 0.0;
  m_track_t = 0.0;

  if (good_track!=NULL) {
    m_track_recon = true;
    m_track_t = good_track->time();

    vector<const DBCALShower*> bcal_showers;
    if (is_charged) {
      good_track->Get(bcal_showers);
    } else {
      //need to go one step deeper in neutral case
      vector<const DNeutralShower*> neutral_showers;
      good_track->Get(neutral_showers);
      if (neutral_showers.size() > 0) neutral_showers[0]->Get(bcal_showers);
    }
    m_bcal_match = bcal_showers.size() > 0; //we never have more than one match
    if (m_bcal_match) {
      m_cluster_t = bcal_showers[0]->t;
      m_cluster_E = bcal_showers[0]->E;
      m_cluster_num_points = bcal_showers[0]->N_cell;
      vector<const DBCALCluster*> bcal_clusters;
      vector<const DBCALPoint*> bcal_points;
      bcal_showers[0]->Get(bcal_clusters);
      if (bcal_clusters.size() == 1) bcal_clusters[0]->Get(bcal_points);
      sort(bcal_points.begin(),bcal_points.end(),TimeSort);
      //Get time of first hit
      if (bcal_points.size() > 0) m_cluster_t_first = bcal_points[0]->t();
      //Get average time of earliest half of hits
      int half_num_points = bcal_points.size()/2; //round down
      double wt_sum = 0.0;
      for (int i=0; i<half_num_points; i++) {
        double wt;
        if (bcal_points[i]->layer() != 4) {
          wt = bcal_points[i]->E()*bcal_points[i]->E();
        } else {
          wt = 0.0;
        }
        wt_sum += wt;
        m_cluster_t_early_half += bcal_points[i]->tInnerRadius()*wt;
      }
      if (wt_sum > 0.0) { 
        m_cluster_t_early_half /= wt_sum;
        //cluster t is at inner radius, correct for this
        m_cluster_t_early_half *= sqrt(bcal_showers[0]->x*bcal_showers[0]->x + bcal_showers[0]->y*bcal_showers[0]->y)/64.3;
      } else {//if we only have one point or if all first half points are in the fourth layer, wt_sum will equal 0. just use the original cluster time in this case.
        m_cluster_t_early_half = bcal_showers[0]->t;

      }
      
      //Get average time from first half of points closest to entry point
      double x = bcal_showers[0]->x;
      double y = bcal_showers[0]->y;
      double r = sqrt(x*x +y*y);
      double z = bcal_showers[0]->z;
      //Coordinates of the entry point
      gl_x = x*64.3/r;
      gl_y = y*64.3/r;
      gl_z = (z-65.)*64.3/r; //the z we will be comparing to, DBCALPoint::z(), is relative to the center of the target so we don't need to add the 65 back in
      sort(bcal_points.begin(),bcal_points.end(),PointToPointDistSort);
      wt_sum = 0.0;
      for (int i=0; i<half_num_points; i++) {
        double wt;
        if (bcal_points[i]->layer() != 4) {
          wt = bcal_points[i]->E()*bcal_points[i]->E();
        } else {
          wt = 0.0;
        }
        wt_sum += wt;
        m_cluster_t_close_to_entry_half += bcal_points[i]->tInnerRadius()*wt;
      }
      if (wt_sum > 0.0) { 
        m_cluster_t_close_to_entry_half /= wt_sum;
        //cluster t is at inner radius, correct for this
        m_cluster_t_close_to_entry_half *= sqrt(bcal_showers[0]->x*bcal_showers[0]->x + bcal_showers[0]->y*bcal_showers[0]->y)/64.3;
      } else {//if we only have one point or if all first half points are in the fourth layer, wt_sum will equal 0. just use the original cluster time in this case.
        m_cluster_t_close_to_entry_half = bcal_showers[0]->t;

      }

      //Get average time from first half of points closest to line from target center to cluster position
      sort(bcal_points.begin(),bcal_points.end(),PointToLineDistSort);
      wt_sum = 0.0;
      for (int i=0; i<half_num_points; i++) {
        double wt;
        if (bcal_points[i]->layer() != 4) {
          wt = bcal_points[i]->E()*bcal_points[i]->E();
        } else {
          wt = 0.0;
        }
        wt_sum += wt;
        m_cluster_t_close_to_track_half += bcal_points[i]->tInnerRadius()*wt;
      }
      if (wt_sum > 0.0) { 
        m_cluster_t_close_to_track_half /= wt_sum;
        //cluster t is at inner radius, correct for this
        m_cluster_t_close_to_track_half *= sqrt(bcal_showers[0]->x*bcal_showers[0]->x + bcal_showers[0]->y*bcal_showers[0]->y)/64.3;
      } else {//if we only have one point or if all first half points are in the fourth layer, wt_sum will equal 0. just use the original cluster time in this case.
        m_cluster_t_close_to_track_half = bcal_showers[0]->t;

      }

      sort(bcal_points.begin(),bcal_points.end(),TimeSort);
      wt_sum = 0.0;
      int max_points = bcal_points.size();
      if (max_points > 1) max_points -= max_points/4 + 1;
      for (int i=0; i<max_points; i++) {

        double wt;
        if (bcal_points[i]->layer() != 4) {
          wt = bcal_points[i]->E()*bcal_points[i]->E();
        } else {
          wt = 0.0;
        }
        wt_sum += wt;
        m_cluster_t_grabbag += bcal_points[i]->tInnerRadius()*wt;
      }
      if (wt_sum > 0.0) { 
        m_cluster_t_grabbag /= wt_sum;
        //cluster t is at inner radius, correct for this
        m_cluster_t_grabbag *= sqrt(bcal_showers[0]->x*bcal_showers[0]->x + bcal_showers[0]->y*bcal_showers[0]->y)/64.3;
      } else {//if we only have one point or if all first half points are in the fourth layer, wt_sum will equal 0. just use the original cluster time in this case.
        m_cluster_t_grabbag = bcal_showers[0]->t;

      }

      //Save info about all points to a different tree
      for (int i=0; i<bcal_points.size(); i++) {
        m_point_E = bcal_points[i]->E();
        m_point_t = bcal_points[i]->t();
        m_point_t_inner = bcal_points[i]->tInnerRadius();
        m_point_z = bcal_points[i]->z();
        m_point_dist_to_entry = DistToEntry(bcal_points[i]);
        m_point_dist_to_track = DistToTrack(bcal_points[i]);
        m_point_layer = bcal_points[i]->layer();
        associated_points_tree->Fill();
      }
    }
  }


  bcal_charged_timing_tree->Fill();

  UnlockState();

  return NOERROR;
}

//------------------
// erun
//------------------
jerror_t DEventProcessor_bcal_charged_timing::erun(void)
{
  // Any final calculations on histograms (like dividing them)
  // should be done here. This may get called more than once.
  return NOERROR;
}

//------------------
// fini
//------------------
jerror_t DEventProcessor_bcal_charged_timing::fini(void)
{
  return NOERROR;
}