#include <IECalPionControlSample.hxx>
#include <IReconTrack.hxx>
#include <IReconPID.hxx>
#include <IAlgorithmResult.hxx>
#include <IRealDatum.hxx>
#include <ICOMETLog.hxx>
#include <ReconObjectUtils.hxx>
#include <TrackingUtils.hxx>
#include <IOARuntimeParameters.hxx>
#include <TMath.h>

IECalPionControlSample::IECalPionControlSample() :
  IControlSampleBase("ECalPion") {

  fMinTpcHits = COMET::IOARuntimeParameters::Get().GetParameterI("oaControlSample.ECalPion.MinTpcHits");

  fMinMuonPullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MinMuonPullAccept");
  fMaxMuonPullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MaxMuonPullAccept");
  fMinProtonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MinProtonPullReject");
  fMaxProtonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MaxProtonPullReject");

  fMinMomentum = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MinMomentum");
  fMaxMomentum = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.ECalPion.MaxMomentum");

  AddValidTrigger(COMET::Trigger::kBeamSpill);
}

IECalPionControlSample::~IECalPionControlSample() {
}

double IECalPionControlSample::GetFrontMomentum(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconState> firstState = TrackingUtils::GetFirstState(*object);

  if (!firstState) {
    return -10000.;
  }

  return TrackingUtils::GetMomentum(firstState);
}

double IECalPionControlSample::GetCharge( COMET::IHandle<COMET::IReconBase> object){
  COMET::IHandle<COMET::IReconState> firstState = TrackingUtils::GetFirstState(*object);

  if (!firstState) {
    return -1000.;
  }
  
  double charge = TrackingUtils::GetCharge(firstState);
  return charge;
}

bool IECalPionControlSample::PassesPullCuts(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconPID> pid = object;
  double pullMuon = 10000, pullProton = 10000;

  if (pid) {
    COMET::IHandle<COMET::IRealDatum> test_datum = (pid)->Get<COMET::IRealDatum> ("tpcPid_NTrun");

    if (test_datum) {
      pullMuon = pid->Get<COMET::IRealDatum> ("tpcPid_PullMuon")->GetValue();
      pullProton = pid->Get<COMET::IRealDatum> ("tpcPid_PullProton")->GetValue();
    } else {
      return false;
    }
  }

  if (pullMuon > fMinMuonPullAccept && pullMuon < fMaxMuonPullAccept) {
    return false;
  }

  if (pullProton > fMinProtonPullReject && pullProton < fMaxProtonPullReject) {
    return false;
  }

  return true;
}

bool IECalPionControlSample::AtECalFrontFace(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconState> lastState = TrackingUtils::GetLastState(*object);

  if (!lastState) {
    return false;
  }

  TLorentzVector backPos = TrackingUtils::GetPosition(lastState);

  double barrelLeft = 890.;
  double barrelRight = -890.;
  double barrelBottom = -980.;
  double barrelTop = 1085.;
  double barrelFront = 600.;
  double barrelBack = 2600.;
  double DsFront    = 2665;
  double DsTop      = 930;
  double DsBottom   = -910;
  double DsLeft     = 920;
  double DsRight    = -920;

  if ((backPos.X() > barrelLeft || backPos.X() < barrelRight || backPos.Y() < barrelBottom || backPos.Y() > barrelTop) && (backPos.Z() > barrelFront && backPos.Z() < barrelBack))
  {
    return true; // entered Barrel
  }

  else if ( backPos.X() < DsLeft && backPos.X() > DsRight && backPos.Y() > DsBottom && backPos.Y() < DsTop && backPos.Z() > DsFront )
  {
    return true; // entered DsECal
  }
  else {return false;}
}

int IECalPionControlSample::FGDStartPoint(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconState> firstState = TrackingUtils::GetFirstState(*object);

  if (!firstState) {
    return false;
  }

  TLorentzVector pos = TrackingUtils::GetPosition(firstState);
  TVector3 fgd1Min = TVector3(-832.2, -777.2, 123.45);
  TVector3 fgd1Max = TVector3(832.2, 887.2, 446.95);
  TVector3 fgd2Min = TVector3(-832.2, -777.2, 1481.45);
  TVector3 fgd2Max = TVector3(832.2, 887.2, 1807.95);

  if (pos.X() < fgd1Max.X() &&
      pos.X() > fgd1Min.X() &&
      pos.Y() < fgd1Max.Y() &&
      pos.Y() > fgd1Min.Y() &&
      pos.Z() < fgd1Max.Z() &&
      pos.Z() > fgd1Min.Z()  ){
    return 1;
    }

  else if (pos.X() < fgd2Max.X() &&
           pos.X() > fgd2Min.X() &&
           pos.Y() < fgd2Max.Y() &&
           pos.Y() > fgd2Min.Y() &&
           pos.Z() < fgd2Max.Z() &&
           pos.Z() > fgd2Min.Z()  ){
       return 2;
       }
  else {return 0;}
}

// method checks which bunch the object is in, returns an integer between 1 and 8
int IECalPionControlSample::WhichBunch(COMET::IHandle<COMET::IReconBase> object, int eventID) {
  COMET::IHandle<COMET::IReconState> firstState = TrackingUtils::GetFirstState(*object);
  
  if (!firstState) {
  return false;
  }
  
  double bunchTimeArray[8];
  if (eventID < 6000){
  bunchTimeArray[0] = 2839.7;  bunchTimeArray[1] = 3423.5;  bunchTimeArray[2] = 4005.4;  bunchTimeArray[3] = 4588.6;  bunchTimeArray[4] = 5172.2;  bunchTimeArray[5] = 5754.6;  bunchTimeArray[6] = -1000;  bunchTimeArray[7] = -1000;
  }
  else if (eventID < 7000){
  bunchTimeArray[0] = 2853.95;  bunchTimeArray[1] = 3444.15;  bunchTimeArray[2] = 4030.41;  bunchTimeArray[3] = 4620.34;  bunchTimeArray[4] = 5180.28;  bunchTimeArray[5] = 5770.12;  bunchTimeArray[6] = 6343.77;  bunchTimeArray[7] = 6924.67;
  }
 else if (eventID < 8000){
  bunchTimeArray[0] = 3019.11;  bunchTimeArray[1] = 3597.74;  bunchTimeArray[2] = 4180.73;  bunchTimeArray[3] = 4763.93;  bunchTimeArray[4] = 5346.49;  bunchTimeArray[5] = 5927.83;  bunchTimeArray[6] = 6508.5;  bunchTimeArray[7] = 7093.56;
  }
  else
  {
  bunchTimeArray[0] = 2752.;  bunchTimeArray[1] = 3333.;  bunchTimeArray[2] = 3915.;  bunchTimeArray[3] = 4498.;  bunchTimeArray[4] = 5080.;  bunchTimeArray[5] = 5661.;  bunchTimeArray[6] = 6244.;  bunchTimeArray[7] = 6826.;
  }

  double time = TrackingUtils::GetPosition(firstState).T();
  double bunchWidth = 75;

  for (int i = 1; i < 9; i++){
    double bunchMin = bunchTimeArray[i] - bunchWidth;
    double bunchMax = bunchTimeArray[i] + bunchWidth;

  if (time > bunchMin && time < bunchMax){
  return i;
  }
  
  }
  return 0;
}


bool IECalPionControlSample::IsEventSelectedInternal(COMET::ICOMETEvent& event) {
  // find all tracks in same FGD, in the same bunch, with good TPC components. If > 1: check all for pull/charge/momentum cuts
  // map contains pair(FGD, bunch), vector(objects)
  std::map< std::pair<int, int> , std::vector< COMET::IHandle<COMET::IReconBase> > > foundObjects;
  std::vector< COMET::IHandle<COMET::IReconBase> >  selectedObjects;

  COMET::IHandle<COMET::IAlgorithmResult> reconResult = event.GetFit("ReconGlobal");
  if (!reconResult) {
    return false;
  }

  COMET::IHandle<COMET::IReconObjectContainer> roc = reconResult->GetResultsContainer("final");
  if (!roc) {
    return false;
  }

  for (COMET::IReconObjectContainer::iterator it = roc->begin(); it != roc->end(); it++) {
    COMET::IHandle<COMET::IReconBase> object = *it;

    if (!object || !object->UsesDetector(COMET::IReconBase::kTPC) || !object->CheckStatus(COMET::IReconBase::kSuccess)) {
      continue;
    }

    COMET::IHandle<COMET::IReconBase> tpcObject;

    if (object->UsesDetector(COMET::IReconBase::kTPC3)) {
      tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC3);
    } else if (object->UsesDetector(COMET::IReconBase::kTPC2)) {
      tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC2);
    } else if (object->UsesDetector(COMET::IReconBase::kTPC1)) {
      tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC1);
    }

    if (!tpcObject) {
      continue;
    }

  int eventID = event.GetEventId();

  int whichBunch = WhichBunch(object, eventID);
  int whichFGD = FGDStartPoint(object);
  
  // ignore object if it doesnt start within a bunch time or in an FGD FV
  if (whichFGD == 0 || whichBunch == 0){
  continue;
  }
    
  int nHits = 0;
    if (tpcObject->GetHits()) {
      nHits = tpcObject->GetHits()->size();
    }
    std::pair<int, int> key (whichFGD, whichBunch);
    if (nHits > fMinTpcHits) {
      foundObjects[key].push_back(tpcObject);
    }
  }
  
  // loop through map to see if any of them have more than one object
  std::map< std::pair<int, int>, std::vector<COMET::IHandle<COMET::IReconBase> > >::iterator it;
  for (it = foundObjects.begin(); it != foundObjects.end(); it++){
    std::vector<COMET::IHandle<COMET::IReconBase> > currentObject = (*it).second;
  
  if (currentObject.size() > 1)
  {
      // loop through found objects - if pull cuts etc passed, fill selectedObjects
    for (std::vector<COMET::IHandle<COMET::IReconBase> >::iterator object_it = currentObject.begin() ; object_it !=  currentObject.end(); object_it++)
    {
      COMET::IHandle<COMET::IReconBase> object = (*object_it);
        COMET::IHandle<COMET::IReconBase> tpcObject;

        if (object->UsesDetector(COMET::IReconBase::kTPC3)) {
          tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC3);
        } else if (object->UsesDetector(COMET::IReconBase::kTPC2)) {
          tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC2);
        } else if (object->UsesDetector(COMET::IReconBase::kTPC1)) {
          tpcObject = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC1);
        }
          
    if (!tpcObject || !AtECalFrontFace(tpcObject)) {
          continue;
        }

        bool passesPulls = PassesPullCuts(tpcObject);
        double momentum = GetFrontMomentum(tpcObject);
        double charge = GetCharge(tpcObject);

    if (passesPulls &&
        momentum > fMinMomentum &&
        momentum < fMaxMomentum &&
      charge > 0
        ) {
         selectedObjects.push_back(object);
    }
    }
    }
  }

  if (selectedObjects.size() == 0){
    return false;
    }

  return true;
}