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

IDsECalProtonControlSample::IDsECalProtonControlSample() :
  IControlSampleBase("DsECalProton") {

  fMinTpc3Hits = COMET::IOARuntimeParameters::Get().GetParameterI("oaControlSample.DsECalProton.MinTpc3Hits");

  fMinProtonPullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MinProtonPullAccept");
  fMinProtonPullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MinProtonPullAccept");
  fMinElePullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MinElePullReject");
  fMaxElePullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MaxElePullReject");
  fMinMuonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MinMuonPullReject");
  fMaxMuonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MaxMuonPullReject");
  fMinPionPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MinPionPullReject");
  fMaxPionPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.DsECalProton.MaxPionPullReject");

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

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

IDsECalProtonControlSample::~IDsECalProtonControlSample() {
  
}


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

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

  return TrackingUtils::GetMomentum(firstState);
}

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

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

  return TrackingUtils::GetCharge(firstState);
}

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

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

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

  if (pullMuon > fMinMuonPullReject && pullMuon < fMaxMuonPullReject) {
    return false;
  }

  if (pullPion > fMinPionPullReject && pullPion < fMaxPionPullReject) {
    return false;
  }

  return true;
}

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

  if (!lastState) {
    return false;
  }

  TLorentzVector backPos = TrackingUtils::GetPosition(lastState);
  double dsECalMinZ = 2665.;
  double dsECalMinX = -920.;
  double dsECalMaxX = 920.;
  double dsECalMinY = -910.;
  double dsECalMaxY = 930.;

  double barrelMaxX = 890.;
  double barrelMinX = -890.;
  double barrelMinY = -980.;
  double barrelMaxY = 1085.;
  double barrelMinZ = 600.;
  double barrelMaxZ = 2600.;

  if (backPos.X() < dsECalMaxX && backPos.X() > dsECalMinX && backPos.Y() < dsECalMaxY && backPos.Y() > dsECalMinY && backPos.Z() > dsECalMinZ)
  {return true;} // entered DsECal
  
  else if ( (backPos.X() > barrelMaxX || backPos.X() < barrelMinX || backPos.Y() > barrelMaxY || backPos.Y() < barrelMinY) && backPos.Z() > barrelMinZ && backPos.Z() < barrelMaxZ)
  {return true;} // entered Barrel

  else {return false;}
}

bool IDsECalProtonControlSample::IsEventSelectedInternal(COMET::ICOMETEvent& event) {
  std::ofstream myfile;
  myfile.open("output/protonEventIDs.dat", std::ios::app);
  COMET::IHandle<COMET::IAlgorithmResult> reconResult = event.GetFit("ReconGlobal");
  if (!reconResult) {
    myfile.close();
    return false;
  }

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


  // The TPC3 objects that enter the DsECal and pass all our other cuts
  std::vector<COMET::IHandle<COMET::IReconBase> > selectedObjects;

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

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

    COMET::IHandle<COMET::IReconBase> tpc3Object = ReconObjectUtils::GetConstituentInDetector(object, COMET::IReconBase::kTPC3);

    if (!tpc3Object || !AtDsECalFrontFace(tpc3Object)) {
      continue;
    }

    bool passesPulls = PassesPullCuts(tpc3Object);
    double momentum = GetFrontMomentum(tpc3Object);
    double charge = GetCharge(tpc3Object);
    int nHits = 0;
    

  if (charge < 0){
    myfile.close();
  return false;
  }

  if (tpc3Object->GetHits()) {
      nHits = tpc3Object->GetHits()->size();
    }

    if (nHits > fMinTpc3Hits && momentum > fMinMomentum && momentum < fMaxMomentum && passesPulls) {
      selectedObjects.push_back(tpc3Object);
    }
  }

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

  int eventID = event.GetEventId();
  int runID = event.GetRunId();

  myfile << eventID << ", " << runID << std::endl;
  myfile.close();
  return true;
}