// PndAnalysis
// Needs a FairRunAna set up in the macro for file & parameter I/O

#include "PndAnalysis.h"

#include <string>
#include <iostream>

using std::cout;
using std::endl;

//Root stuff
#include "TTree.h"
#include "TChain.h"
#include "TClonesArray.h"
#include "TParticle.h"
#include "TDatabasePDG.h"
#include "TParticlePDG.h"

#include "VAbsPidSelector.h"
#include "VAbsMicroCandidate.h"

//RHO stuff
#include "TRho.h"
#include "TCandidate.h"
#include "TCandList.h"
#include "TPidSelector.h"

#include "PndPidProbability.h"
#include "PndPidListMaker.h"
#include "PndMCTrack.h"

ClassImp(PndAnalysis);

PndAnalysis::PndAnalysis() :
fRootManager(FairRootManager::Instance()),
fPidListMaker(0),
fEvtCount(0),
fChainEntries(0),
fEventRead(false),
fBuildMcCands(false),
fVerbose(0),
fChargedPidName("PidAlgoIdealCharged"),
fNeutralPidName("PidAlgoIdealNeutral")
{
  if ( 0 == fRootManager )
  {
    std::cout << "-E- PndAnalysis: RootManager not instantiated!" << std::endl;
    return;
  }
  Init();
}

PndAnalysis::~PndAnalysis()
{
	if(0!=fPidListMaker) delete fPidListMaker;
}

TClonesArray* PndAnalysis::ReadTCA(TString tcaname)
{
  TClonesArray* tca = (TClonesArray*) fRootManager->GetObject(tcaname.Data());
  if (! tca) std::cout << "-W- PndAnalysis::Init(): No "<<tcaname.Data()<<" array found." << std::endl;
  return tca;
}

void PndAnalysis::Init()
{
	Reset();
  
  //read arrays
  fChargedCands = ReadTCA("PidChargedCand");  
  fChargedProbability = ReadTCA(fChargedPidName.Data());
  fNeutralCands = ReadTCA("PidNeutralCand");
  fNeutralProbability = ReadTCA(fChargedPidName.Data());
  fBuildMcCands = false;
  fMcCands = ReadTCA("PndMcTracks");
  if ( ! fMcCands )
  {
    if( fVerbose ) std::cout << "-I- PndAnalysis::Init(): Trying mc stack now." << std::endl;
    fMcTracks = (TClonesArray*) fRootManager->GetObject("MCTrack");
    if ( ! fMcTracks && fVerbose ) std::cout << "-W- PndAnalysis::Init(): No \"MCTrack\" array found. No MC info available." << std::endl;
    fMcCands =new TClonesArray("TCandidate");
    fRootManager->Register("PndMcTracks","PndMcTracks", fMcCands, kTRUE);
    fBuildMcCands = true;
  }
  fChainEntries =(fRootManager->GetInChain())->GetEntries();
  
	fPidListMaker = new PndPidListMaker();
  
  fPdg = TRho::Instance()->GetPDG();
  
}


void PndAnalysis::Rewind()
{
	fEvtCount=0;
}

Int_t PndAnalysis::GetEvent(Int_t n)
{	
	allCands.Cleanup();
	chargedCands.Cleanup();
	neutralCands.Cleanup();
	mcCands.Cleanup();
	
	fEventRead=false;
	
	if (n>=0) fEvtCount=n+1;
	else fEvtCount++;
	
	if (fEvtCount<=fChainEntries) return fEvtCount;
	else fEvtCount=fChainEntries;
	
	return 0;
}

Bool_t PndAnalysis::FillList(TCandList &l, std::string listkey)
{
  // Reads the specified List for the current event
  
	l.Cleanup();
	
	
	// when the first list is requested read in the event
	if (!fEventRead) 
	{
    fRootManager->ReadEvent(fEvtCount-1);
		fEventRead=true;
	}
  
	if (listkey=="McTruth")
	{
		if (fMcCands) {
      if(fBuildMcCands) BuildMcCands();
      for (Int_t i1=0; i1<fMcCands->GetEntriesFast(); i1++){
        TCandidate* tc = (TCandidate *)fMcCands->At(i1);
        l.Add(*tc);
      }
      return true;
    } else return false;
	}
	
  if (allCands.GetLength() == 0) // do only when we didn't read something yet.
	{ 	// removed now compatibility to TCandidate readin ... instead read PndPidCandidates
		if (fNeutralCands && neutralCands.GetLength()==0)
      for (Int_t i1=0; i1<fNeutralCands->GetEntriesFast(); i1++)
      {
        VAbsMicroCandidate *mic = (VAbsMicroCandidate *)fNeutralCands->At(i1);		
        TCandidate tc(*mic,i1+1);
        // TODO: Do we want to set something here? It is neutrals anyway.
        if(i1<fNeutralProbability->GetEntriesFast())
        {
          PndPidProbability *neuProb = (PndPidProbability*)fNeutralProbability->At(i1);
          // numbering see PndPidListMaker
          tc.SetPidInfo(0,neuProb->GetElectronPidProb());
          tc.SetPidInfo(1,neuProb->GetMuonPidProb());
          tc.SetPidInfo(2,neuProb->GetPionPidProb());
          tc.SetPidInfo(3,neuProb->GetKaonPidProb());
          tc.SetPidInfo(4,neuProb->GetProtonPidProb());
        }        
        neutralCands.Add(tc);
        allCands.Add(tc);
      }
    
		if (fChargedCands && chargedCands.GetLength()==0) 
    {
      for (Int_t i1=0; i1<fChargedCands->GetEntriesFast(); i1++)
      {
        VAbsMicroCandidate *mic = (VAbsMicroCandidate *)fChargedCands->At(i1);
        TCandidate tc(*mic,i1+1);
        if(i1<fChargedProbability->GetEntriesFast())
        {
          PndPidProbability *chProb = (PndPidProbability*)fChargedProbability->At(i1);
          // numbering see PndPidListMaker
          tc.SetPidInfo(0,chProb->GetElectronPidProb());
          tc.SetPidInfo(1,chProb->GetMuonPidProb());
          tc.SetPidInfo(2,chProb->GetPionPidProb());
          tc.SetPidInfo(3,chProb->GetKaonPidProb());
          tc.SetPidInfo(4,chProb->GetProtonPidProb());
        }        
        chargedCands.Add(tc);
        allCands.Add(tc);
      }
    }
	}
	
	
	// set the base list for the PID list maker
	fPidListMaker->SetBaseList(chargedCands);
	
	if (listkey=="All" )
	{
		l=allCands;
		return true;
	}
  
	if (listkey=="Neutral") 
	{
		l=neutralCands;
		return true;
	}
	
	if (listkey=="Charged") 
	{
		l=chargedCands;
		return true;
	}
  
	return fPidListMaker->FillList(l,listkey);
}

Int_t PndAnalysis::GetEntries() 
{
  if (fRootManager) return (fRootManager->GetInChain())->GetEntries();
	else return 0;
}

void PndAnalysis::BuildMcCands()
{
  if (fMcCands->GetEntriesFast() != 0) fMcCands->Delete();
  //if (fMcCands->GetEntriesFast() != 0)  fMcCands->Clear("C");
  
  // Get the Candidates
  for(Int_t i=0; i<fMcTracks->GetEntriesFast(); i++)
  {
  	PndMCTrack *part = (PndMCTrack*)fMcTracks->At(i);
  	if (part->GetMotherID()!=-1) continue;
  	
    TLorentzVector p4 = part->Get4Momentum();
    TVector3    stvtx = part->GetStartVertex();
    
    TParticlePDG *ppdg = fPdg->GetParticle(part->GetPdgCode());
    
    double charge=0.0;
    
    if (ppdg) charge=ppdg->Charge();
    else if (fVerbose) cout <<"-W- PndMcListConverter: strange PDG code:"<<part->GetPdgCode()<<endl;
    if (fabs(charge)>2) charge/=3.;
    
    //TClonesArray& ref = *fMcCands;
    Int_t size = fMcCands->GetEntriesFast();
    
   	TCandidate *pmc=new ((*fMcCands)[size]) TCandidate(p4,charge);
   	
    pmc->SetMcIdx(size);
    pmc->SetPos(stvtx);
    pmc->SetType(part->GetPdgCode());
  }
  
  if(fVerbose) cout <<"-I- PndMcListConverter: found primaries="<<fMcCands->GetEntriesFast()<<endl;
  
}