//
// Pibero Djawotho <pibero@indiana.edu>
// Indiana University
// 14 August 2007
//

// ROOT
#include "TFile.h"
#include "TH1.h"
#include "TH2.h"
#include "TNtuple.h"
#include "TSystem.h"

// STAR
#include "StEventTypes.h"
#include "StMcEvent/StMcEventTypes.hh"
#include "StAssociationMaker/StAssociationMaker.h"
#include "StAssociationMaker/StTrackPairInfo.hh"
#include "StEventUtilities/StuRefMult.hh"

// Local
#include "StMcUpsilonQAMaker.h"

ClassImp(StMcUpsilonQAMaker);

int StMcUpsilonQAMaker::Init()
{
  const char* filename = "ups.mc.root";
  if (gSystem->Getenv("JOBID"))
    filename = gSystem->ConcatFileName(gSystem->Getenv("SCRATCH"), filename);
  mFile = new TFile(filename, "recreate");
  assert(mFile && !mFile->IsZombie());

  TString varlist;
  varlist += "vxMc:vyMc:vzMc:cosThetaMc";
  varlist += ":upsilonId:upsilonMomMc:upsilonPtMc:upsilonPzMc:upsilonRapidityMc:upsilonEtaMc:upsilonPhiMc";
  varlist += ":electronId:electronMomMc:electronPtMc:electronPzMc:electronRapidityMc:electronEtaMc:electronPhiMc";
  varlist += ":positronId:positronMomMc:positronPtMc:positronPzMc:positronRapidityMc:positronEtaMc:positronPhiMc";
  mMcTuple = new TNtuple("mc", "Upsilon MC Tuple", varlist);

  TString varlist2 = varlist;
  varlist2 += ":refmult:cosThetaRc";
  varlist2 += ":upsilonMassRc:upsilonMomRc:upsilonPtRc:upsilonPzRc:upsilonRapidityRc:upsilonEtaRc:upsilonPhiRc";
  varlist2 += ":electronMomRc:electronPtRc:electronPzRc:electronRapidityRc:electronEtaRc:electronPhiRc:electronFitPoints:electronDca";
  varlist2 += ":positronMomRc:positronPtRc:positronPzRc:positronRapidityRc:positronEtaRc:positronPhiRc:positronFitPoints:positronDca";
  mRcTuple = new TNtuple("rc", "Upsilon Reco Tuple", varlist2);

  mUpsilonPtMc  = new TH1F("upsilonPtMc" , ";p_{T} [GeV/c]", 10, 0, 5);
  mUpsilonPtRc  = new TH1F("upsilonPtRc" , ";p_{T} [GeV/c]", 10, 0, 5);
  mUpsilonPtEff = new TH1F("upsilonPtEff", ";p_{T} [GeV/c]", 10, 0, 5);

  mElectronPtMc  = new TH1F("electronPtMc" , ";p_{T} [GeV/c]", 20, 0, 10);
  mElectronPtRc  = new TH1F("electronPtRc" , ";p_{T} [GeV/c]", 20, 0, 10);
  mElectronPtEff = new TH1F("electronPtEff", ";p_{T} [GeV/c]", 20, 0, 10);

  mPositronPtMc  = new TH1F("positronPtMc" , ";p_{T} [GeV/c]", 20, 0, 10);
  mPositronPtRc  = new TH1F("positronPtRc" , ";p_{T} [GeV/c]", 20, 0, 10);
  mPositronPtEff = new TH1F("positronPtEff", ";p_{T} [GeV/c]", 20, 0, 10);

  mUpsilonPtMc->Sumw2();
  mUpsilonPtRc->Sumw2();
  mUpsilonPtEff->Sumw2();

  mElectronPtMc->Sumw2();
  mElectronPtRc->Sumw2();
  mElectronPtEff->Sumw2();

  mPositronPtMc->Sumw2();
  mPositronPtRc->Sumw2();
  mPositronPtEff->Sumw2();

  return StMaker::Init();
}

int StMcUpsilonQAMaker::Make()
{
  // Load StMcEvent
  StMcEvent* mcEvent = (StMcEvent*)GetDataSet("StMcEvent");
  if (!mcEvent) {
    LOG_WARN << "No StMcEvent" << endm;
    return kStWarn;
  }

  // Skip event with no Monte Carlo primary vertex
  if (!mcEvent->primaryVertex()) return kStOk;

  // Track loop
  for (unsigned int i = 0; i < mcEvent->primaryVertex()->numberOfDaughters(); ++i) {
    StMcTrack* mcTrack = mcEvent->primaryVertex()->daughter(i);
    if (mcTrack->geantId() == 160 || mcTrack->geantId() == 161 ||
	mcTrack->geantId() == 162 || mcTrack->geantId() == 163) {

      if (mcTrack->stopVertex() && mcTrack->stopVertex()->numberOfDaughters() == 2) {
	StMcTrack* mcTrack1 = mcTrack->stopVertex()->daughter(0);
	StMcTrack* mcTrack2 = mcTrack->stopVertex()->daughter(1);

	if (mcTrack1->geantId() != 3) swap(mcTrack1, mcTrack2);

	float cosTheta = mcTrack1->momentum().unit() * mcTrack2->momentum().unit();

	float tuple[] = {
	  mcEvent->primaryVertex()->position().x(),
	  mcEvent->primaryVertex()->position().y(),
	  mcEvent->primaryVertex()->position().z(),
	  cosTheta,

	  // Upsilon
	  mcTrack->geantId(),
	  mcTrack->momentum().mag(),
	  mcTrack->pt(),
	  mcTrack->momentum().z(),
	  mcTrack->rapidity(),
	  mcTrack->pseudoRapidity(),
	  mcTrack->momentum().phi(),

	  // Electron
	  mcTrack1->geantId(),
	  mcTrack1->momentum().mag(),
	  mcTrack1->pt(),
	  mcTrack1->momentum().z(),
	  mcTrack1->rapidity(),
	  mcTrack1->pseudoRapidity(),
	  mcTrack1->momentum().phi(),

	  // Positron
	  mcTrack2->geantId(),
	  mcTrack2->momentum().mag(),
	  mcTrack2->pt(),
	  mcTrack2->momentum().z(),
	  mcTrack2->rapidity(),
	  mcTrack2->pseudoRapidity(),
	  mcTrack2->momentum().phi()
	};

	mMcTuple->Fill(tuple);
      }
    }
  }

  // Load StEvent
  StEvent* event = (StEvent*)GetDataSet("StEvent");
  if (!event) {
    LOG_WARN << "No StEvent" << endm;
    return kStWarn;
  }

  // Skip event with no rcnstructed primary vertex
  if (!event->primaryVertex()) return kStOk;

  // Reference multiplicity
  int refmult = uncorrectedNumberOfPrimaries(*event);

  static const float eMass = StElectron::instance()->mass();

  // Track loop
  for (unsigned int i = 0; i < mcEvent->primaryVertex()->numberOfDaughters(); ++i) {
    StMcTrack* mcTrack = mcEvent->primaryVertex()->daughter(i);
    if (mcTrack->geantId() == 160 || mcTrack->geantId() == 161 ||
	mcTrack->geantId() == 162 || mcTrack->geantId() == 163) {

      if (mcTrack->stopVertex() && mcTrack->stopVertex()->numberOfDaughters() == 2) {
	StMcTrack* mcTrack1 = mcTrack->stopVertex()->daughter(0);
	StMcTrack* mcTrack2 = mcTrack->stopVertex()->daughter(1);

	if (!(accept(mcTrack1) && accept(mcTrack2))) continue;

	mUpsilonPtMc->Fill(mcTrack->pt());
	mElectronPtMc->Fill(mcTrack1->pt());
	mPositronPtMc->Fill(mcTrack2->pt());

	if (mcTrack1->geantId() != 3) swap(mcTrack1, mcTrack2);

	StTrack* track1 = findPartner(mcTrack1);
	StTrack* track2 = findPartner(mcTrack2);

	if (!(track1 && accept(track1))) continue;
	if (!(track2 && accept(track2))) continue;

	float cosTheta = mcTrack1->momentum().unit() * mcTrack2->momentum().unit();

	mElectronPtRc->Fill(track1->geometry()->momentum().perp());
	mPositronPtRc->Fill(track2->geometry()->momentum().perp());

	// Rcnstruct parent 4-momentum
	const StThreeVectorF& mom1 = track1->geometry()->momentum();
	const StThreeVectorF& mom2 = track2->geometry()->momentum();

	StLorentzVectorF p1(mom1, mom1.massHypothesis(eMass));
	StLorentzVectorF p2(mom2, mom2.massHypothesis(eMass));
	StLorentzVectorF p = p1 + p2;

	mUpsilonPtRc->Fill(p.perp());

	float cosThetaRc = mom1.unit() * mom2.unit();

	float tuple[] = {
	  mcEvent->primaryVertex()->position().x(),
	  mcEvent->primaryVertex()->position().y(),
	  mcEvent->primaryVertex()->position().z(),
	  cosTheta,

	  // Upsilon
	  mcTrack->geantId(),
	  mcTrack->momentum().mag(),
	  mcTrack->pt(),
	  mcTrack->momentum().z(),
	  mcTrack->rapidity(),
	  mcTrack->pseudoRapidity(),
	  mcTrack->momentum().phi(),

	  // Electron
	  mcTrack1->geantId(),
	  mcTrack1->momentum().mag(),
	  mcTrack1->pt(),
	  mcTrack1->momentum().z(),
	  mcTrack1->rapidity(),
	  mcTrack1->pseudoRapidity(),
	  mcTrack1->momentum().phi(),

	  // Positron
	  mcTrack2->geantId(),
	  mcTrack2->momentum().mag(),
	  mcTrack2->pt(),
	  mcTrack2->momentum().z(),
	  mcTrack2->rapidity(),
	  mcTrack2->pseudoRapidity(),
	  mcTrack2->momentum().phi(),

	  refmult,
	  cosThetaRc,

	  // Upsilon
	  p.m(),
	  p.vect().mag(),
	  p.perp(),
	  p.pz(),
	  p.rapidity(),
	  p.pseudoRapidity(),
	  p.phi(),

	  // Electron
	  p1.vect().mag(),
	  p1.perp(),
	  p1.pz(),
	  p1.rapidity(),
	  p1.pseudoRapidity(),
	  p1.phi(),
	  track1->fitTraits().numberOfFitPoints(),
	  dca(track1->node()->track(global), event->primaryVertex()->position()).mag(),

	  // Positron
	  p2.vect().mag(),
	  p2.perp(),
	  p2.pz(),
	  p2.rapidity(),
	  p2.pseudoRapidity(),
	  p2.phi(),
	  track2->fitTraits().numberOfFitPoints(),
	  dca(track2->node()->track(global), event->primaryVertex()->position()).mag()
	};

	mRcTuple->Fill(tuple);
      }
    }
  }

  return kStOk;
}

int StMcUpsilonQAMaker::Finish()
{
  mUpsilonPtEff ->Divide(mUpsilonPtRc , mUpsilonPtMc );
  mElectronPtEff->Divide(mElectronPtRc, mElectronPtMc);
  mPositronPtEff->Divide(mPositronPtRc, mPositronPtMc);
  mFile->Write();
  mFile->Close();
  return kStOk;
}

bool StMcUpsilonQAMaker::accept(StMcTrack* mcTrack) const
{
  return !mcTrack->bemcHits().empty();
}

bool StMcUpsilonQAMaker::accept(StTrack* track) const
{
  return (track && track->flag() > 0 &&
	  track->fitTraits().numberOfFitPoints() > 20 &&
	  track->fitTraits().numberOfFitPoints() / (float)track->numberOfPossiblePoints() > 0.52 &&
	  track->geometry()->momentum().perp() > 0.2 &&
	  track->geometry()->momentum().mag() < 100);
}

StTrack* StMcUpsilonQAMaker::findPartner(StMcTrack* mcTrack)
{
  StAssociationMaker* assoc = (StAssociationMaker*)GetMaker("StAssociationMaker");
  pair<mcTrackMapIter, mcTrackMapIter> p = assoc->mcTrackMap()->equal_range(mcTrack);
  StTrack* maxTrack = 0;
  int maxCommonTpcHits = 0;
  for (mcTrackMapIter k = p.first; k != p.second; ++k) {
    int commonTpcHits = k->second->commonTpcHits();
    StTrack* track = k->second->partnerTrack()->node()->track(primary);
    if (track && commonTpcHits > maxCommonTpcHits) {
      maxTrack = track;
      maxCommonTpcHits = commonTpcHits;
    }
  }
  return maxTrack;
}

StThreeVectorF StMcUpsilonQAMaker::dca(StTrack* track, const StThreeVectorF& point) const
{
  StPhysicalHelixD helix = track->geometry()->helix();
  double pathlength = helix.pathLength(point, false); // Do not scan periods
  return helix.at(pathlength) - point;
}
