#include <IBarrelECalElectronControlSample.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>

IBarrelECalElectronControlSample::IBarrelECalElectronControlSample() :
  IControlSampleBase("BarrelECalElectron") {

  fMinTpcHits = COMET::IOARuntimeParameters::Get().GetParameterI("oaControlSample.BarrelECalElectron.MinTpcHits");
  fTpcBackDirMinAngle = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.TpcBackDirMinAngle");
  fTpcBackDirMaxAzimuth = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.TpcBackDirMaxAzimuth");
  fTpcBackDirMinAzimuth = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.TpcBackDirMinAzimuth");

  fMinElePullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MinElePullAccept");
  fMaxElePullAccept = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MaxElePullAccept");
  fMinMuonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MinMuonPullReject");
  fMaxMuonPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MaxMuonPullReject");
  fMinPionPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MinPionPullReject");
  fMaxPionPullReject = COMET::IOARuntimeParameters::Get().GetParameterD("oaControlSample.BarrelECalElectron.MaxPionPullReject");

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

  fMustStartInFGD = COMET::IOARuntimeParameters::Get().GetParameterI("oaControlSample.BarrelECalElectron.MustStartInFGD");

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

IBarrelECalElectronControlSample::~IBarrelECalElectronControlSample() {
  if (hMom) {
    delete hMom;
  }
}

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

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

  return TrackingUtils::GetMomentum(firstState);
}

double IBarrelECalElectronControlSample::GetBackAngleToZ(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconState> lastState = TrackingUtils::GetLastState(*object);

  if (!lastState) {
    return -1000.;
  }

  return TrackingUtils::GetDirection(lastState).Unit().Angle(TVector3(0, 0, 1)) * TMath::RadToDeg();
}

double IBarrelECalElectronControlSample::GetBackAzimuth(COMET::IHandle<COMET::IReconBase> object) {
  COMET::IHandle<COMET::IReconState> lastState = TrackingUtils::GetLastState(*object);

  if (!lastState) {
    return -1000.;
  }

  TVector3 tpcBackDir = TrackingUtils::GetDirection(lastState);
  double azimuth = TMath::RadToDeg() * TMath::ATan(tpcBackDir.X() / tpcBackDir.Y());

  if (tpcBackDir.X() > 0 && tpcBackDir.Y() < 0) {
    azimuth += 180;
  } else if (tpcBackDir.X() < 0 && tpcBackDir.Y() < 0) {
    azimuth -= 180;
  }

  return azimuth;
}

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

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

    if (test_datum) {
      pullEle = pid->Get<COMET::IRealDatum> ("tpcPid_PullEle")->GetValue();
      pullMuon = pid->Get<COMET::IRealDatum> ("tpcPid_PullMuon")->GetValue();
      pullPion = pid->Get<COMET::IRealDatum> ("tpcPid_PullPion")->GetValue();
    } else {
      return false;
    }
  }

  if (pullEle < fMinElePullAccept || pullEle > fMaxElePullAccept) {
    return false;
  }

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

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

  return true;
}

bool IBarrelECalElectronControlSample::AtBarrelECalFrontFace(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.;

  return ((backPos.X() > barrelLeft || backPos.X() < barrelRight || backPos.Y() < barrelBottom || backPos.Y() > barrelTop) && (backPos.Z() > barrelFront && backPos.Z() < barrelBack));
}

bool IBarrelECalElectronControlSample::StartsInFGD(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);

  return ((pos.X() < fgd1Max.X() &&
           pos.X() > fgd1Min.X() &&
           pos.Y() < fgd1Max.Y() &&
           pos.Y() > fgd1Min.Y() &&
           pos.Z() < fgd1Max.Z() &&
           pos.Z() > fgd1Min.Z()    ) ||
          (pos.X() < fgd2Max.X() &&
           pos.X() > fgd2Min.X() &&
           pos.Y() < fgd2Max.Y() &&
           pos.Y() > fgd2Min.Y() &&
           pos.Z() < fgd2Max.Z() &&
           pos.Z() > fgd2Min.Z()    )    );
}

void IBarrelECalElectronControlSample::InitHistograms() {
  hMom = new TH1F("hMom", "Momentum", 50, 0, 5000);
}

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

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

  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::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 || !AtBarrelECalFrontFace(tpcObject)) {
      continue;
    }

    bool inFGD = (fMustStartInFGD == 0 || StartsInFGD(object));
    bool passesPulls = PassesPullCuts(tpcObject);
    double momentum = GetFrontMomentum(tpcObject);
    double backAngle = GetBackAngleToZ(tpcObject);
    double backAzimuth = GetBackAzimuth(tpcObject);

    int nHits = 0;
    if (tpcObject->GetHits()) {
      nHits = tpcObject->GetHits()->size();
    }

    //std::cout << "nHits " << nHits << ", backAngle " << backAngle << ", backAzimuth " << backAzimuth << ", momentum " << momentum << ", passesPulls " << passesPulls << std::endl;

    if (nHits > fMinTpcHits &&
        backAngle > fTpcBackDirMinAngle &&
        backAzimuth > fTpcBackDirMinAzimuth &&
        backAzimuth < fTpcBackDirMaxAzimuth &&
        momentum > fMinMomentum &&
        momentum < fMaxMomentum &&
        passesPulls &&
        inFGD) {
      selectedObjects.push_back(tpcObject);
    }
  }

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

  if(fMakeHistoFile)
    for (std::vector<COMET::IHandle<COMET::IReconBase> >::iterator it = selectedObjects.begin(); it != selectedObjects.end(); it++) {
      hMom->Fill(GetFrontMomentum(*it));
    }

  return true;
}