#include "modules/PythiaFix.h"


#include "ExRootAnalysis/ExRootResult.h"
#include "ExRootAnalysis/ExRootClasses.h"

#include "ExRootAnalysis/ExRootFilter.h"
#include "ExRootAnalysis/ExRootClassifier.h"

#include "ExRootAnalysis/ExRootFactory.h"
#include "ExRootAnalysis/ExRootCandidate.h"

#include "TMath.h"
#include "TString.h"
#include "TLorentzVector.h"
#include "TDatabasePDG.h"
#include "TClonesArray.h"

#include <iostream>
#include <set>

using namespace std;

//------------------------------------------------------------------------------

PythiaFix::PythiaFix() :
  fItBranchParticle(0)
{
}

//------------------------------------------------------------------------------

PythiaFix::~PythiaFix()
{
}

//------------------------------------------------------------------------------

void PythiaFix::Init()
{
  // import ROOT tree branch

  fBranchParticle = UseBranch("GenParticle");
  fItBranchParticle = fBranchParticle->MakeIterator();

}

//------------------------------------------------------------------------------

void PythiaFix::Finish()
{ 
  if(fItBranchParticle) delete fItBranchParticle;
}

//------------------------------------------------------------------------------

void PythiaFix::FixDaughters(Int_t indexParticle, Int_t indexMother)
{
  ExRootGenParticle *mother = 0;
  mother = static_cast<ExRootGenParticle*>(fBranchParticle->At(indexMother));

  if(!mother) return;

  // skip beam particles
  if(mother->M1 < 0 && mother->M2 < 0) return;

  if(mother->D1 < 0 && mother->D2 < 0)
  {
    mother->D1 = indexParticle;
		mother->D2 = indexParticle;
  }
  else if(mother->D1 < indexParticle && mother->D2 + 1 == indexParticle)
  {
		mother->D2 = indexParticle;
  }
  else if(mother->D1 < indexParticle)
  {
    mother->D1 = indexParticle;
		mother->D2 = indexParticle;
  }
}

//------------------------------------------------------------------------------

void PythiaFix::Process()
{
  ExRootGenParticle *particle = 0;
  Int_t indexParticle, indexMother;

/*
  TDatabasePDG *pdg = TDatabasePDG::Instance();

  indexParticle = -1;
  fItBranchParticle->Reset();
  while(particle = static_cast<ExRootGenParticle*>(fItBranchParticle->Next()))
  {
    ++indexParticle;
    cout << indexParticle << "\t" << pdg->GetParticle(particle->PID)->GetName() << "\t";
    cout << particle->Status << "\t" << particle->M1 << "\t" << particle->M2 << "\t";
    cout << particle->D1 << "\t" << particle->D2 << endl;
  }
*/

  indexParticle = -1;
  fItBranchParticle->Reset();
  while(particle = static_cast<ExRootGenParticle*>(fItBranchParticle->Next()))
  {
		++indexParticle;

		if(particle->M1 < 0) continue;
		
		if(particle->M2 < 0)
		{
      FixDaughters(indexParticle, particle->M1);
    }
    else
    {
      for(indexMother = particle->M1; indexMother <= particle->M2; ++indexMother)
      {
        FixDaughters(indexParticle, indexMother);
      }
    }
  }

/*
  cout << "==============================" << endl;
  indexParticle = -1;
  fItBranchParticle->Reset();
  while(particle = static_cast<ExRootGenParticle*>(fItBranchParticle->Next()))
  {
    ++indexParticle;
    cout << indexParticle << "\t" << pdg->GetParticle(particle->PID)->GetName() << "\t";
    cout << particle->Status << "\t" << particle->M1 << "\t" << particle->M2 << "\t";
    cout << particle->D1 << "\t" << particle->D2 << endl;
  }
*/
}

//------------------------------------------------------------------------------