#include "DEventProcessor_bcal_tree.h"

#include <TLorentzVector.h>

#include "DANA/DApplication.h"
#include "BCAL/DBCALShower.h"
#include "BCAL/DBCALTruthShower.h"
#include "BCAL/DBCALCluster.h"
#include "BCAL/DBCALGeometry.h"
#include "BCAL/DBCALTruthCell.h"
#include "BCAL/DBCALIncidentParticle.h"
#include "FCAL/DFCALCluster.h"
#include "TRACKING/DMCThrown.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_tree());
  }
} // "C"


//------------------
// init
//------------------
jerror_t DEventProcessor_bcal_tree::init(void)
{
  // Create BCAL directory
  TDirectory *dir = new TDirectoryFile("BCAL","BCAL");
  dir->cd();

  //set up the tree
  BCAL_shower_info = new TTree("BCAL_shower_info","BCAL shower info");

  BCAL_shower_info->Branch("E_thrown",&m_E_thrown,"E_thrown/F");
  BCAL_shower_info->Branch("phi_thrown",&m_phi_thrown,"phi_thrown/F");
  BCAL_shower_info->Branch("theta_thrown",&m_theta_thrown,"theta_thrown/F");
  BCAL_shower_info->Branch("z_entry_thrown",&m_z_entry_thrown,"z_entry_thrown/F");
  BCAL_shower_info->Branch("vertex_z_thrown",&m_vertex_z_thrown,"vertex_z_thrown/F");

  BCAL_shower_info->Branch("nShowers",&m_nShowers,"nShowers/I");

  BCAL_shower_info->Branch("match",&m_match,"match/O");
  BCAL_shower_info->Branch("matched_z_entry",&m_matched_z_entry,"matched_z_entry/F");
  BCAL_shower_info->Branch("matched_E_raw",&m_matched_E_raw,"matched_E_raw/F");
  BCAL_shower_info->Branch("matched_E",&m_matched_E,"matched_E/F");
  BCAL_shower_info->Branch("matched_sig_E",&m_matched_sig_E,"matched_sig_E/F");
  BCAL_shower_info->Branch("matched_t0",&m_matched_t0,"matched_t0/F");
  BCAL_shower_info->Branch("matched_sig_t",&m_matched_sig_t,"matched_sig_t/F");
  BCAL_shower_info->Branch("matched_x",&m_matched_x,"matched_x/F");
  BCAL_shower_info->Branch("matched_y",&m_matched_y,"matched_y/F");
  BCAL_shower_info->Branch("matched_z",&m_matched_z,"matched_z/F");
  BCAL_shower_info->Branch("matched_phi",&m_matched_phi,"matched_phi/F");
  BCAL_shower_info->Branch("matched_sig_phi",&m_matched_sig_phi,"matched_sig_phi/F");
  BCAL_shower_info->Branch("matched_theta",&m_matched_theta,"matched_theta/F");
  BCAL_shower_info->Branch("matched_sig_theta",&m_matched_sig_theta,"matched_sig_theta/F");
  BCAL_shower_info->Branch("matched_nCells",&m_matched_nCells,"matched_nCells/I");

  BCAL_shower_info->Branch("no_conversions",&m_no_conversions,"no_conversions/O");
  BCAL_shower_info->Branch("no_conversions_incident_particle",&m_no_conversions_incident_particle,"no_conversions_incident_particle/O");

  BCAL_shower_info->Branch("nPoints",&m_nPoints,"nPoints/I");
  BCAL_shower_info->Branch("E_truthcell",m_E_truthcell,"E_truthcell[nPoints]/F");
  BCAL_shower_info->Branch("t_truthcell",m_t_truthcell,"t_truthcell[nPoints]/F");
  BCAL_shower_info->Branch("z_truthcell",m_z_truthcell,"z_truthcell[nPoints]/F");
  BCAL_shower_info->Branch("E_point",m_E_point,"E_point[nPoints]/F");
  BCAL_shower_info->Branch("t_point",m_t_point,"t_point[nPoints]/F");
  BCAL_shower_info->Branch("z_point",m_z_point,"z_point[nPoints]/F");
  
  BCAL_shower_info->Branch("E_hit_up",m_E_hit_up,"E_hit_up[nPoints]/F");
  BCAL_shower_info->Branch("E_hit_down",m_E_hit_down,"E_hit_down[nPoints]/F");

  // Go back up to the parent directory
  dir->cd("../");

  return NOERROR;
}

//------------------
// brun
//------------------
jerror_t DEventProcessor_bcal_tree::brun(JEventLoop *eventLoop, int runnumber)
{
  return NOERROR;
}

//------------------
// evnt
//------------------
jerror_t DEventProcessor_bcal_tree::evnt(JEventLoop *loop, int eventnumber)
{
  vector<const DBCALShower*> showers;
  vector<const DMCThrown*> mcthrowns;
  vector<const DBCALTruthShower*> showers_truth;
  vector<const DBCALIncidentParticle*> incident_particles;
  loop->Get(showers);
  loop->Get(mcthrowns);
  loop->Get(showers_truth);
  loop->Get(incident_particles);
  
  LockState();

  /* Section about comparing DBCALTruthCell (pre-smeared info) to DBCALPoint */
  vector<const DBCALTruthCell*> truth_cells;
  vector<const DBCALPoint*> points;
  loop->Get(truth_cells);
  loop->Get(points);

  m_nPoints=0;
  for (vector<const DBCALPoint*>::iterator points_i = points.begin();
       points_i != points.end();
       ++points_i) {

    if (m_nPoints >= MAX_BCALPOINTS) continue;

    double truthcell_E=0.0, truthcell_t=0.0, truthcell_z=0.0;
    int truthcell_Ncells=0;
    int point_cellId = DBCALGeometry::cellId( (**points_i).module(), (**points_i).layer(), (**points_i).sector() );
    for (vector<const DBCALTruthCell*>::iterator cells_i = truth_cells.begin();
         cells_i != truth_cells.end();
         ++cells_i) {
      int fADC_cellId = DBCALGeometry::fADCId( (**cells_i).module, (**cells_i).layer, (**cells_i).sector );
      if (point_cellId == fADC_cellId) {
        truthcell_E += (**cells_i).E;
        truthcell_t += (**cells_i).t * (**cells_i).E;
        truthcell_z += (**cells_i).zLocal * (**cells_i).E;
        truthcell_Ncells++;
      }
    }

    truthcell_t /= truthcell_E;
    truthcell_z /= truthcell_E;
    m_E_truthcell[m_nPoints] = truthcell_E;
    m_t_truthcell[m_nPoints] = truthcell_t;
    m_z_truthcell[m_nPoints] = truthcell_z + DBCALGeometry::GLOBAL_CENTER;
    m_E_point[m_nPoints] = (**points_i).E();
    m_t_point[m_nPoints] = (**points_i).t();
    m_z_point[m_nPoints] = (**points_i).z() + 65.0;

    
    vector<const DBCALHit*> hits;
    (**points_i).Get(hits);
    if (hits.size()==2) {
      for (int i=0; i<2; i++) {
        if (hits[i]->end == DBCALGeometry::kUpstream) {
          m_E_hit_up[m_nPoints] = hits[i]->E;
        } else if (hits[i]->end == DBCALGeometry::kDownstream) 
          m_E_hit_down[m_nPoints] = hits[i]->E;{
        }
      }
    }

    m_nPoints++;
  }
  
  /* End of section comparing DBCALTruthCell to DBCALPoint*/

  //reset everything so there are no leftovers from the last event
  m_E_thrown=0;
  m_phi_thrown=0;
  m_theta_thrown=0;
  m_z_entry_thrown=0;
  m_vertex_z_thrown=0;

  m_match=false;
  m_matched_z_entry=0;
  m_matched_E_raw=0;
  m_matched_E=0;
  m_matched_t0=0;
  m_matched_x=0;
  m_matched_y=0;
  m_matched_z=0;
  m_matched_phi=0;
  m_matched_sig_phi=0;
  m_matched_theta=0;
  m_matched_sig_theta=0;
  m_matched_nCells=0;

  m_no_conversions=false;
  m_no_conversions_incident_particle=false;

  //is there a truth shower coming from a primary particle? This means that the photon reached the BCAL without converting
  m_no_conversions = (showers_truth.size()>=1) && showers_truth[0]->primary;
  //Try to get the equivalent information from DBCALIncidentParticle, seems to give slightly different results
  m_no_conversions_incident_particle = (incident_particles.size()==1) && incident_particles[0]->ptype==1;

  const double pi=TMath::Pi();
  //very rarely we end up with >1 DMCThrown (due to interactions outside
  //the target?). I am not really sure why, but it should not be a big deal;
  //we can still assume the first DMCThrown is the original thrown photon.
  if (mcthrowns.size()>=1) {
    DVector3 vertex = mcthrowns[0]->position();
    DLorentzVector momentum = mcthrowns[0]->lorentzMomentum();

    m_E_thrown=momentum.E();
    m_phi_thrown=momentum.Phi();
    m_theta_thrown=momentum.Theta();

    m_vertex_z_thrown = vertex.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...
    const double bcal_inner_r = DBCALGeometry::BCALINNERRAD;
    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);
    m_z_entry_thrown = s*momentum.Z()+vertex.Z();
  }

  m_nShowers=showers.size();

  m_match=false; //is there a shower that corresponds to thrown particle (in phi)
  for (unsigned int i=0; i<showers.size();i++) {
    //as of svn revision 9763 the t stored is the time at the shower
    //position (x,y,z) rather than the time at the inner radius
    //with this change, we have consistency across algorithms

    //get associated cluster,as this gives us useful objects like sig_phi
    vector<const DBCALCluster*> cluster;
    showers[i]->Get(cluster);
    double sig_phi = 0;
    double sig_theta = 0;
    double nCells = 0;
    if (cluster.size()==1) {
      sig_phi = cluster[0]->sigPhi();
      sig_theta = cluster[0]->sigTheta();
      nCells = cluster[0]->nCells();
    }

    double x = showers[i]->x;
    double y = showers[i]->y;
    double inner_rad = DBCALGeometry::BCALINNERRAD;
    double z = showers[i]->z;
    double t = showers[i]->t;
    double sig_t = showers[i]->tErr;
    double E = showers[i]->E;
    double sig_E = 0.0564*sqrt(E) + 0.0064*E; //this is how it is computed in DNeutralShowerCandidate.h
    double E_raw = showers[i]->E_raw;
    double r = sqrt( x*x + y*y );
    //since we are using this z_entry to determine if a shower matches the thrown photon we will use the
    //thrown vertex position
    double z_entry = (z-m_vertex_z_thrown)*inner_rad/r + m_vertex_z_thrown;
    //For calculating t0 (e.g. delta_TOF), we won't assume that we have that we
    //perfect vertex resolution, so just z=65 as target position
    double path_full = sqrt( r*r + (z-65.)*(z-65.)); //should not be hard-coded
    double t0_full = t - (path_full/(30*k_cm/k_nsec));

    double phi = atan2(y,x);
    //use z=65 as target position here
    double theta = atan2(sqrt(x*x+y*y),z-65.);

    double deltaPhi = phi-m_phi_thrown;
    double deltaPhiAlt = ( phi > m_phi_thrown ?
			   phi - m_phi_thrown - 2*PI :
			   m_phi_thrown - phi - 2*PI );     
    deltaPhi = min( fabs( deltaPhi ), fabs( deltaPhiAlt ) );
    double delta_phi_thresh = .07;

    double delta_E = E - m_E_thrown;
    double sig_E_NIM = m_E_thrown*sqrt(.054*.054/m_E_thrown + .023*.023); //nominal energy resolution from NIM paper
    double delta_E_thresh = min(3.0*sig_E_NIM,1.0*m_E_thrown);
    
    double delta_z = z_entry - m_z_entry_thrown;
    double sig_z_NIM = 1.1/sqrt(m_E_thrown); //nominal z-resolution from NIM paper
    double delta_z_thresh = min(5.*sig_z_NIM,39.0);
    delta_z_thresh = max(delta_z_thresh, 9.0);

    if (deltaPhi < delta_phi_thresh && fabs(delta_z) < delta_z_thresh && fabs(delta_E) < delta_E_thresh) {
      if (m_match==true) cerr << ">1 correspondence" << endl;
      m_match=true;
      m_matched_z_entry = z_entry;
      m_matched_E_raw = E_raw;
      m_matched_E = E;
      m_matched_sig_E = sig_E;
      m_matched_t0 = t0_full;
      m_matched_sig_t = sig_t;
      m_matched_x = x;
      m_matched_y = y;
      m_matched_z = z;
      m_matched_phi = phi;
      m_matched_sig_phi = sig_phi;
      m_matched_theta = theta;
      m_matched_sig_theta = sig_theta;
      m_matched_nCells = nCells;

    }
    phi = phi<0 ? 2*pi+phi : phi;
  }
  cout << boolalpha << m_match << endl;

  BCAL_shower_info->Fill();

  UnlockState();

  return NOERROR;
}

//------------------
// erun
//------------------
jerror_t DEventProcessor_bcal_tree::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_tree::fini(void)
{
  return NOERROR;
}