#include <IBasketTrackSandMuonControlSample.hxx>
#include <IReconTrack.hxx>
#include <IReconPID.hxx>
#include "CLHEP/Matrix/Matrix.h"
#include <IFieldManager.hxx>
#include <IHelixModel.hxx>

#define TPC2MMMINZ 595.0
#define TPC2MMMAXZ 1320.0

#define TPCMMMAXY 950.0
#define TPCMMMINY -850.0
#define TPCMMMAXX 790.0
#define TPCMMMINX -790.0


//*************************************************************
void TTrackStateToHelix( COMET::IHandle<COMET::ITrackState> state, IHelixModel *Helix){
//*************************************************************
  double x0 = state->GetPosition().X();
  double y0 = state->GetPosition().Y();
  double z0 = state->GetPosition().Z();

  double tx0 = state->GetDirection().X();
  double ty0 = state->GetDirection().Y()/ state->GetDirection().Z();

  double rho = state->GetCurvature();

  Helix->SetParameters(x0,y0,z0,tx0,ty0,rho);
}

//*************************************************************
void GetMomentum( COMET::IHandle<COMET::ITrackState> state,double &p,double &ep, double &pt, double &ept){
//*************************************************************

  // need to read in magnet field...
  double B =  COMET::IFieldManager::GetFieldValue( state->GetPosition() ).Mag()/unit::tesla;

  // should get speed of light constant ~0.3 from somewhere
  // Curvature in 1/mm
  // p in MeV/c
  p = 0.3*B/(state->GetCurvature()*TMath::Sqrt(1.-state->GetDirection()[0]*state->GetDirection()[0]));

  pt = 0.3*B/state->GetCurvature();

  CLHEP::HepMatrix  Jacobian(2,2,0);
  CLHEP::HepMatrix  Covariance(2,2,0);
  CLHEP::HepMatrix  Cov(2,2,0);
  
  Covariance[0][0] = state->GetCovarianceValue(state->GetDirectionIndex(),state->GetDirectionIndex());
  Covariance[0][1] = state->GetCovarianceValue(state->GetDirectionIndex(),state->GetCurvatureIndex());

  Covariance[1][0] = state->GetCovarianceValue(state->GetCurvatureIndex(),state->GetDirectionIndex());
  Covariance[1][1] = state->GetCovarianceValue(state->GetCurvatureIndex(),state->GetCurvatureIndex());

  ept =  0.3*B/state->GetCurvature()/state->GetCurvature()*TMath::Sqrt(Covariance[1][1]);
  
  Jacobian[0][0] = 1.;
  Jacobian[0][1] = 0.;

  Jacobian[1][0] = p/(1.-state->GetDirection()[0]*state->GetDirection()[0])*TMath::Abs(state->GetDirection()[0]);
  Jacobian[1][1] = -p/state->GetCurvature(); 

  Cov = Jacobian*Covariance*Jacobian.T();
 
  ep = TMath::Sqrt(Cov[1][1]);

  p = TMath::Abs(p);
 
  return ; 
}



IBasketTrackSandMuonControlSample::IBasketTrackSandMuonControlSample() : IControlSampleBase("BasketTrackSandMuon"){

  AddValidTrigger(COMET::Trigger::kBeamSpill);

}

void IBasketTrackSandMuonControlSample::InitHistograms(){

}

bool IBasketTrackSandMuonControlSample::IsEventSelectedInternal(COMET::ICOMETEvent& event) {

  // First check the P0D
  const TVector3& kP0DMaxPos=COMET::IGeomInfo::P0D().ActiveMax();
  TVector3 P0DLastPos(-0xABCDEF, -0xABCDEF, -0xABCDEF);
  int NbTracksP0D=0;
  
  for(int i = 0; i < 24; i++){
    COMET::IHandle<COMET::IAlgorithmResult> algres = 
      event.GetFit(TString::Format("ReconP0D/p0dCycle%d/IP0DTrackRecon", i));
    if (!algres){
      continue;
    }
    // COMETLog("Got fit in cycle" << i);
    COMET::IHandle<COMET::IReconObjectContainer> vtxs=algres->GetResultsContainer();
    // Require exactly one vertex
    if(!vtxs || vtxs->size()!=1) continue;
    
    COMET::IHandle<COMET::IReconVertex> vtx=vtxs->front();
    if(!vtx) COMETWarn("Something bad happened - no IReconVertex");
    
    
    // Count tracks in vertex
    COMET::IHandle<COMET::IReconObjectContainer> constituents=vtx->GetConstituents();
    
    for(int j=0; j < (int)constituents->size(); ++j){
      COMET::IHandle<COMET::IReconPID> trk=constituents->at(j);
      if(trk){
        COMET::IHandle<COMET::IHitSelection> trkHits=trk->GetHits();
        
        const TVector3& frontPos=trkHits->front()->GetPosition();
        const TVector3& backPos=trkHits->back()->GetPosition();
        if(fabs(frontPos.Z()-backPos.Z()) > 72*unit::cm && 
           fabs(backPos.Z()-kP0DMaxPos.Z()) < 2*unit::cm){
          NbTracksP0D++;
          P0DLastPos = backPos;
        }
      }
    }
  }
    
  // Then check the number of tracks in the TPCs
  COMET::IHandle<COMET::IAlgorithmResult> Result = event.GetFit("ReconTPC");
  if ( ! Result ) 
    return false;
  
  int NbTracksTPC1 = 0;
  int NbTracksTPC2 = 0;
  int NbTracksTPC3 = 0;
  
  bool TPC1toTPC2 = false;
  bool TPC3toTPC2 = false;
  bool TPC2toP0D  = false;
  
  COMET::IHandle<COMET::IReconObjectContainer> tpcTracks = Result->GetResultsContainer("TPCRecontracks");
  if ( ! tpcTracks )
    return false;
  
  for ( COMET::IReconObjectContainer::iterator it=tpcTracks->begin(); it!=tpcTracks->end(); it++ ){
    COMET::IHandle< COMET::IReconTrack > Track = *it;
    COMET::IHandle<COMET::ITrackState> State = Track->GetState();
    double Momentum, eMomentum;
    double pt, ept; 
    GetMomentum(State, Momentum, eMomentum,pt,ept);
    // Don't count the delta rays
    if (Momentum < 50. )
      continue;
    
    IHelixModel *Helix = new IHelixModel();
    TTrackStateToHelix(State, Helix);

    if ( (*it)->UsesDetector(COMET::IReconBase::kTPC1)){
      NbTracksTPC1++;
      double Xprop = Helix->XCoord(TVector3(0.,0.,TPC2MMMAXZ),2);
      double Yprop = Helix->YCoord(TVector3(0.,0.,TPC2MMMAXZ),2);
      TPC1toTPC2 = Xprop < TPCMMMAXX && Xprop > TPCMMMINX && Yprop < TPCMMMAXY && Yprop > TPCMMMINY;
    } else if ( (*it)->UsesDetector(COMET::IReconBase::kTPC2)){
      NbTracksTPC2++;
      double XP0DTPC2Dist = P0DLastPos.X() - Helix->XCoord(P0DLastPos,2);
      double YP0DTPC2Dist = P0DLastPos.Y() - Helix->YCoord(P0DLastPos,2);
      TPC2toP0D = fabs(XP0DTPC2Dist) < 200. && fabs(YP0DTPC2Dist) < 200.;
    } else if ( (*it)->UsesDetector(COMET::IReconBase::kTPC3)){
      NbTracksTPC3++;
      double Xprop = Helix->XCoord(TVector3(0.,0.,TPC2MMMINZ),2);
      double Yprop = Helix->YCoord(TVector3(0.,0.,TPC2MMMINZ),2);
      TPC3toTPC2 = Xprop < TPCMMMAXX && Xprop > TPCMMMINX && Yprop < TPCMMMAXY && Yprop > TPCMMMINY;
    }
  }
  
  
  if (NbTracksTPC1== 1 && NbTracksTPC3== 1 && TPC1toTPC2 && TPC3toTPC2)
    return true;
  if (NbTracksP0D == 1 && NbTracksTPC2== 1 && TPC2toP0D)
    return true;
  
  return false;  

}