//////////////////////////////////////////////////////
//
// PID performance checker of the TRD
//
// Performs checks of ESD information for TRD-PID and recalculate PID based on OCDB information
// Also provides performance plots on detector based on the PID information - for the moment only 
// MC source is used but V0 is also possible. Here is a list of detector performance checks
//   - Integrated dE/dx per chamber
//   - <PH> as function of time bin and local radial position
//   - number of clusters/tracklet 
//   - number of tracklets/track 
//
// Author : Alex Wilk <wilka@uni-muenster.de>
//          Alex Bercuci <A.Bercuci@gsi.de>
//          Markus Fasel <M.Fasel@gsi.de>
//
///////////////////////////////////////////////////////

#include "TAxis.h"
#include "TROOT.h"
#include "TPDGCode.h"
#include "TCanvas.h"
#include "TF1.h"
#include "TH1F.h"
#include "TH1D.h"
#include "TH2F.h"
#include "TProfile.h"
#include "TProfile2D.h"
#include "TGraph.h"
#include "TGraphErrors.h"
#include "TLegend.h"

#include <TClonesArray.h>
#include <TObjArray.h>
#include <TList.h>

#include "AliESDEvent.h"
#include "AliESDInputHandler.h"
#include "AliTrackReference.h"

#include "AliAnalysisTask.h"

#include "AliTRDtrackerV1.h"
#include "AliTRDtrackV1.h"
#include "AliTRDcluster.h"
#include "AliTRDReconstructor.h"
#include "AliCDBManager.h"
#include "AliTRDpidUtil.h"

#include "AliTRDcheckPID.h"
#include "info/AliTRDtrackInfo.h"

ClassImp(AliTRDcheckPID)

//________________________________________________________________________
AliTRDcheckPID::AliTRDcheckPID() 
  :AliTRDrecoTask("checkPID", "TRD PID checker")
  ,fReconstructor(0x0)
  ,fUtil(0x0)
  ,fGraph(0x0)
  ,fEfficiency(0x0)
  ,fMomentumAxis(0x0)
  ,fMinNTracklets(AliTRDgeometry::kNlayer)
  ,fMaxNTracklets(AliTRDgeometry::kNlayer)
 {
  //
  // Default constructor
  //

  fReconstructor = new AliTRDReconstructor();
  fReconstructor->SetRecoParam(AliTRDrecoParam::GetLowFluxParam());

  // Initialize momentum axis with default values
  Double_t defaultMomenta[AliTRDCalPID::kNMom+1];
  for(Int_t imom = 0; imom < AliTRDCalPID::kNMom+1; imom++)
    defaultMomenta[imom] = AliTRDCalPID::GetMomentumBinning(imom);
  SetMomentumBinning(AliTRDCalPID::kNMom, defaultMomenta);

  fUtil = new AliTRDpidUtil();
  InitFunctorList();
}


//________________________________________________________________________
AliTRDcheckPID::~AliTRDcheckPID() 
{
 if(fGraph){fGraph->Delete(); delete fGraph;}
 if(fReconstructor) delete fReconstructor;
 if(fUtil) delete fUtil;
}


//________________________________________________________________________
void AliTRDcheckPID::CreateOutputObjects()
{
  // Create histograms
  // Called once

  OpenFile(0, "RECREATE");
  fContainer = Histos();
}


//_______________________________________________________
TObjArray * AliTRDcheckPID::Histos(){

  //
  // Create QA histograms
  //
  if(fContainer) return fContainer;

  Int_t xBins = AliPID::kSPECIES*fMomentumAxis->GetNbins(); 
  fContainer = new TObjArray(); fContainer->Expand(kNPlots);

  const Float_t epsilon = 1./(2*(AliTRDpidUtil::kBins-1));     // get nice histos with bin center at 0 and 1

  // histos of the electron probability of all 5 particle species and 11 momenta for the 2-dim LQ method 
  fEfficiency = new TObjArray(3);
  fEfficiency->SetOwner(); fEfficiency->SetName("Efficiency");
  fContainer->AddAt(fEfficiency, kEfficiency);
  
  TH1 *h = 0x0;
  if(!(h = (TH2F*)gROOT->FindObject("PID_LQ"))){
    h = new TH2F("PID_LQ", "", xBins, -0.5, xBins - 0.5,
      AliTRDpidUtil::kBins, 0.-epsilon, 1.+epsilon);
  } else h->Reset();
  fEfficiency->AddAt(h, AliTRDpidUtil::kLQ);

  // histos of the electron probability of all 5 particle species and 11 momenta for the neural network method
  if(!(h = (TH2F*)gROOT->FindObject("PID_NN"))){
    h = new TH2F("PID_NN", "", 
      xBins, -0.5, xBins - 0.5,
      AliTRDpidUtil::kBins, 0.-epsilon, 1.+epsilon);
  } else h->Reset();
  fEfficiency->AddAt(h, AliTRDpidUtil::kNN);

  // histos of the electron probability of all 5 particle species and 11 momenta for the ESD output
  if(!(h = (TH2F*)gROOT->FindObject("PID_ESD"))){
    h = new TH2F("PID_ESD", "", 
      xBins, -0.5, xBins - 0.5,
      AliTRDpidUtil::kBins, 0.-epsilon, 1.+epsilon);
  } else h->Reset();
  fEfficiency->AddAt(h, AliTRDpidUtil::kESD);

  // histos of the dE/dx distribution for all 5 particle species and 11 momenta 
  if(!(h = (TH2F*)gROOT->FindObject("dEdx"))){
    h = new TH2F("dEdx", "", 
      xBins, -0.5, xBins - 0.5,
      200, 0, 10000);
  } else h->Reset();
  fContainer->AddAt(h, kdEdx);

  // histos of the dE/dx slices for all 5 particle species and 11 momenta 
  if(!(h = (TH2F*)gROOT->FindObject("dEdxSlice"))){
    h = new TH2F("dEdxSlice", "", 
      xBins*AliTRDpidUtil::kLQslices, -0.5, xBins*AliTRDpidUtil::kLQslices - 0.5,
      200, 0, 5000);
  } else h->Reset();
  fContainer->AddAt(h, kdEdxSlice);

  // histos of the pulse height distribution for all 5 particle species and 11 momenta 
  TObjArray *fPH = new TObjArray(2);
  fPH->SetOwner(); fPH->SetName("PH");
  fContainer->AddAt(fPH, kPH);
  if(!(h = (TProfile2D*)gROOT->FindObject("PHT"))){
    h = new TProfile2D("PHT", "", 
      xBins, -0.5, xBins - 0.5,
      AliTRDtrackerV1::GetNTimeBins(), -0.5, AliTRDtrackerV1::GetNTimeBins() - 0.5);
  } else h->Reset();
  fPH->AddAt(h, 0);
  if(!(h = (TProfile2D*)gROOT->FindObject("PHX"))){
    h = new TProfile2D("PHX", "", 
      xBins, -0.5, xBins - 0.5,
      AliTRDtrackerV1::GetNTimeBins(), 0., .5*AliTRDgeometry::CamHght()+AliTRDgeometry::CdrHght());
  } else h->Reset();
  fPH->AddAt(h, 1);

  // histos of the number of clusters distribution for all 5 particle species and 11 momenta 
  if(!(h = (TH2F*)gROOT->FindObject("NClus"))){
    h = new TH2F("NClus", "", 
      xBins, -0.5, xBins - 0.5,
      50, -0.5, 49.5);
  } else h->Reset();
  fContainer->AddAt(h, kNClus);


  // momentum distributions - absolute and in momentum bins
  if(!(h = (TH1F*)gROOT->FindObject("hMom"))){
    h = new TH1F("hMom", "momentum distribution", fMomentumAxis->GetNbins(), fMomentumAxis->GetXmin(), fMomentumAxis->GetXmax());
  } else h->Reset();
  fContainer->AddAt(h, kMomentum);
  
  if(!(h = (TH1F*)gROOT->FindObject("hMomBin"))){
    h = new TH1F("hMomBin", "momentum distribution in momentum bins", fMomentumAxis->GetNbins(), fMomentumAxis->GetXmin(), fMomentumAxis->GetXmax());
  } else h->Reset();
  fContainer->AddAt(h, kMomentumBin);

  // Number of tracklets per track
  if(!(h = (TH2F*)gROOT->FindObject("nTracklets"))){
    h = new TH2F("nTracklets", "", 
      xBins, -0.5, xBins - 0.5,
      AliTRDgeometry::kNlayer, 0.5, AliTRDgeometry::kNlayer+.5);
  } else h->Reset();
  fContainer->AddAt(h, kNTracklets);

  return fContainer;
}


//________________________________________________________________________
Bool_t AliTRDcheckPID::CheckTrackQuality(const AliTRDtrackV1* track) const
{
  //
  // Check if the track is ok for PID
  //
  
  Int_t ntracklets = track->GetNumberOfTracklets();
  if(ntracklets >= fMinNTracklets && ntracklets <= fMaxNTracklets) return 1;
//   if(!fESD)
//     return 0;

  return 0;
}

//________________________________________________________________________
Int_t AliTRDcheckPID::CalcPDG(AliTRDtrackV1* track) 
{
// Documentation to come

  track -> SetReconstructor(fReconstructor);
  (const_cast<AliTRDrecoParam*>(fReconstructor->GetRecoParam()))->SetPIDNeuralNetwork();
  track -> CookPID();

  if(track -> GetPID(AliPID::kElectron) > track -> GetPID(AliPID::kMuon) + track -> GetPID(AliPID::kPion)  + track -> GetPID(AliPID::kKaon) + track -> GetPID(AliPID::kProton)){
    return kElectron;
  }
  else if(track -> GetPID(kProton) > track -> GetPID(AliPID::kPion)  && track -> GetPID(AliPID::kProton) > track -> GetPID(AliPID::kKaon)  && track -> GetPID(AliPID::kProton) > track -> GetPID(AliPID::kMuon)){
    return kProton;
  }
  else if(track -> GetPID(AliPID::kKaon) > track -> GetPID(AliPID::kMuon)  && track -> GetPID(AliPID::kKaon) > track -> GetPID(AliPID::kPion)){
    return kKPlus;
  }
  else if(track -> GetPID(AliPID::kMuon) > track -> GetPID(AliPID::kPion)){
    return kMuonPlus;
  }
  else{
    return kPiPlus;
  }
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotLQ(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(!fkESD){
    AliWarning("No ESD info available.");
    return 0x0;
  }

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }

  ULong_t status;
  status = fkESD -> GetStatus();
  if(!(status&AliESDtrack::kTRDin)) return 0x0;

  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  if(!(fEfficiency = dynamic_cast<TObjArray *>(fContainer->At(kEfficiency)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }
  TH2F *hPIDLQ = 0x0;
  if(!(hPIDLQ = dynamic_cast<TH2F *>(fEfficiency->At(AliTRDpidUtil::kLQ)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }

  AliTRDtrackV1 cTrack(*fkTrack);
  cTrack.SetReconstructor(fReconstructor);

  Int_t pdg = 0;
  Float_t momentum = 0.;

  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else{
    //AliWarning("No MC info available!");
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;

  (const_cast<AliTRDrecoParam*>(fReconstructor->GetRecoParam()))->SetPIDNeuralNetwork(kFALSE);
  cTrack.CookPID();
  if(cTrack.GetNumberOfTrackletsPID() < fMinNTracklets) return 0x0;

  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  hPIDLQ -> Fill(FindBin(species, momentum), cTrack.GetPID(AliPID::kElectron));
  return hPIDLQ;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotNN(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(!fkESD){
    AliWarning("No ESD info available.");
    return 0x0;
  }

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  ULong_t status;
  status = fkESD -> GetStatus();
  if(!(status&AliESDtrack::kTRDin)) return 0x0;

  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  if(!(fEfficiency = dynamic_cast<TObjArray *>(fContainer->At(kEfficiency)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }
  TH2F *hPIDNN;
  if(!(hPIDNN = dynamic_cast<TH2F *>(fEfficiency->At(AliTRDpidUtil::kNN)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }

  AliTRDtrackV1 cTrack(*fkTrack);
  cTrack.SetReconstructor(fReconstructor);

  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;

  (const_cast<AliTRDrecoParam*>(fReconstructor->GetRecoParam()))->SetPIDNeuralNetwork();
  cTrack.CookPID();
  if(cTrack.GetNumberOfTrackletsPID() < fMinNTracklets) return 0x0;

  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  hPIDNN -> Fill(FindBin(species, momentum), cTrack.GetPID(AliPID::kElectron));
  return hPIDNN;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotESD(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(!fkESD){
    AliWarning("No ESD info available.");
    return 0x0;
  }

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  ULong_t status;
  status = fkESD -> GetStatus();
  if(!(status&AliESDtrack::kTRDin)) return 0x0;

  if(!CheckTrackQuality(fkTrack)) return 0x0;
  if(fkTrack->GetNumberOfTrackletsPID() < fMinNTracklets) return 0x0;
  
  if(!(fEfficiency = dynamic_cast<TObjArray *>(fContainer->At(kEfficiency)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }
  TH2F *hPIDESD = 0x0;
  if(!(hPIDESD = dynamic_cast<TH2F *>(fEfficiency->At(AliTRDpidUtil::kESD)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }


  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    AliTRDtrackV1 cTrack(*fkTrack);
    cTrack.SetReconstructor(fReconstructor);
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;

//   Double32_t pidESD[AliPID::kSPECIES];
  const Double32_t *pidESD = fkESD->GetResponseIter();
  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  hPIDESD -> Fill(FindBin(species, momentum), pidESD[0]);
  return hPIDESD;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotdEdx(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TH2F *hdEdx;
  if(!(hdEdx = dynamic_cast<TH2F *>(fContainer->At(kdEdx)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }

  AliTRDtrackV1 cTrack(*fkTrack);
  cTrack.SetReconstructor(fReconstructor);
  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  if(!IsInRange(momentum)) return 0x0;

  (const_cast<AliTRDrecoParam*>(fReconstructor->GetRecoParam()))->SetPIDNeuralNetwork(kFALSE);
  Float_t sumdEdx = 0;
  Int_t iBin = FindBin(species, momentum);
  AliTRDseedV1 *tracklet = 0x0;
  for(Int_t iChamb = 0; iChamb < AliTRDgeometry::kNlayer; iChamb++){
    sumdEdx = 0;
    tracklet = cTrack.GetTracklet(iChamb);
    if(!tracklet) continue;
    tracklet -> CookdEdx(AliTRDpidUtil::kLQslices);
    for(Int_t i = 3; i--;) sumdEdx += tracklet->GetdEdx()[i];
    hdEdx -> Fill(iBin, sumdEdx);
  }
  return hdEdx;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotdEdxSlice(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TH2F *hdEdxSlice;
  if(!(hdEdxSlice = dynamic_cast<TH2F *>(fContainer->At(kdEdxSlice)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }

  AliTRDtrackV1 cTrack(*fkTrack);
  cTrack.SetReconstructor(fReconstructor);
  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;

  (const_cast<AliTRDrecoParam*>(fReconstructor->GetRecoParam()))->SetPIDNeuralNetwork(kFALSE);
  Int_t iMomBin = fMomentumAxis->FindBin(momentum);
  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  Float_t *fdEdx;
  AliTRDseedV1 *tracklet = 0x0;
  for(Int_t iChamb = 0; iChamb < AliTRDgeometry::kNlayer; iChamb++){
    tracklet = cTrack.GetTracklet(iChamb);
    if(!tracklet) continue;
    tracklet -> CookdEdx(AliTRDpidUtil::kLQslices);
    fdEdx = const_cast<Float_t *>(tracklet->GetdEdx());
    for(Int_t iSlice = 0; iSlice < AliTRDpidUtil::kLQslices; iSlice++){
      hdEdxSlice -> Fill(species * fMomentumAxis->GetNbins() * AliTRDpidUtil::kLQslices + (iMomBin-1) * AliTRDpidUtil::kLQslices + iSlice, fdEdx[iSlice]);
    }
  }  

  return hdEdxSlice;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotPH(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TObjArray *arr = 0x0;
  TProfile2D *hPHX, *hPHT;
  if(!(arr = dynamic_cast<TObjArray *>(fContainer->At(kPH)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }
  hPHT = (TProfile2D*)arr->At(0);
  hPHX = (TProfile2D*)arr->At(1);

  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    AliTRDtrackV1 cTrack(*fkTrack);
    cTrack.SetReconstructor(fReconstructor);
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;;

  AliTRDseedV1 *tracklet = 0x0;
  AliTRDcluster *cluster = 0x0;
  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  Int_t iBin = FindBin(species, momentum);
  for(Int_t iChamb = 0; iChamb < AliTRDgeometry::kNlayer; iChamb++){
    tracklet = fkTrack->GetTracklet(iChamb);
    if(!tracklet) continue;
    Float_t x0 = tracklet->GetX0(); 
    for(Int_t ic = 0; ic < AliTRDtrackerV1::GetNTimeBins(); ic++){
      if(!(cluster = tracklet->GetClusters(ic))) continue;
      hPHT -> Fill(iBin, cluster->GetLocalTimeBin(), TMath::Abs(cluster->GetQ()));
      hPHX -> Fill(iBin, x0 - cluster->GetX(), tracklet->GetdQdl(ic));
    }
  }
  return hPHT;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotNClus(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TH2F *hNClus;
  if(!(hNClus = dynamic_cast<TH2F *>(fContainer->At(kNClus)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }


  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    AliTRDtrackV1 cTrack(*fkTrack);
    cTrack.SetReconstructor(fReconstructor);
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(!IsInRange(momentum)) return 0x0;

  Int_t species = AliTRDpidUtil::AliTRDpidUtil::Pdg2Pid(pdg);
  Int_t iBin = FindBin(species, momentum); 
  AliTRDseedV1 *tracklet = 0x0;
  for(Int_t iChamb = 0; iChamb < AliTRDgeometry::kNlayer; iChamb++){
    tracklet = fkTrack->GetTracklet(iChamb);
    if(!tracklet) continue;
    hNClus -> Fill(iBin, tracklet->GetN());
  }

  return hNClus;
}

//_______________________________________________________
TH1 *AliTRDcheckPID::PlotNTracklets(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  TH2F *hTracklets;
  if(!(hTracklets = dynamic_cast<TH2F *>(fContainer->At(kNTracklets)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }

  AliTRDtrackV1 cTrack(*fkTrack);
  cTrack.SetReconstructor(fReconstructor);
  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  Int_t species = AliTRDpidUtil::Pdg2Pid(pdg);
  if(!IsInRange(momentum)) return 0x0;

  Int_t iBin = FindBin(species, momentum);
  hTracklets -> Fill(iBin, cTrack.GetNumberOfTracklets());
  return hTracklets;
}

//_______________________________________________________
TH1 *AliTRDcheckPID::PlotMom(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TH1F *hMom;
  if(!(hMom = dynamic_cast<TH1F *>(fContainer->At(kMomentum)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }


  Int_t pdg = 0;
  Float_t momentum = 0.;
  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    AliTRDtrackV1 cTrack(*fkTrack);
    cTrack.SetReconstructor(fReconstructor);
    momentum = cTrack.GetMomentum(0);
    pdg = CalcPDG(&cTrack);
  }
  if(IsInRange(momentum)) hMom -> Fill(momentum);
  return hMom;
}


//_______________________________________________________
TH1 *AliTRDcheckPID::PlotMomBin(const AliTRDtrackV1 *track)
{
  //
  // Plot the probabilities for electrons using 2-dim LQ
  //

  if(track) fkTrack = track;
  if(!fkTrack){
    AliWarning("No Track defined.");
    return 0x0;
  }
  
  if(!CheckTrackQuality(fkTrack)) return 0x0;
  
  TH1F *hMomBin;
  if(!(hMomBin = dynamic_cast<TH1F *>(fContainer->At(kMomentumBin)))){
    AliWarning("No Histogram defined.");
    return 0x0;
  }


  Int_t pdg = 0;
  Float_t momentum = 0.;

  if(fkMC){
    if(fkMC->GetTrackRef()) momentum = fkMC->GetTrackRef()->P();
    pdg = fkMC->GetPDG();
  } else {
    //AliWarning("No MC info available!");
    AliTRDtrackV1 cTrack(*fkTrack);
    cTrack.SetReconstructor(fReconstructor);
    momentum = cTrack.GetMomentum(0);
  }
  if(IsInRange(momentum)) hMomBin -> Fill(fMomentumAxis->FindBin(momentum));
  return hMomBin;
}


//________________________________________________________
Bool_t AliTRDcheckPID::GetRefFigure(Int_t ifig)
{
// Steering function to retrieve performance plots

  Bool_t kFIRST = kTRUE;
  TGraphErrors *g = 0x0;
  TAxis *ax = 0x0;
  TObjArray *arr = 0x0;
  TH1 *h1 = 0x0, *h=0x0;
  TH2 *h2 = 0x0;
  TList *content = 0x0;
  switch(ifig){
  case kEfficiency:{
    gPad->Divide(2, 1, 1.e-5, 1.e-5);
    TList *l=gPad->GetListOfPrimitives();
    TVirtualPad *pad = ((TVirtualPad*)l->At(0));pad->cd();
    pad->SetMargin(0.1, 0.01, 0.1, 0.01);

    TLegend *legEff = new TLegend(.65, .75, .98, .98);
    legEff->SetBorderSize(1);
    legEff->SetFillColor(0);
    h=new TH1S("hEff", "", 1, .5, 11.);
    h->SetLineColor(1);h->SetLineWidth(1);
    ax = h->GetXaxis();
    ax->SetTitle("p [GeV/c]");
    ax->SetRangeUser(.5, 11.);
    ax->SetMoreLogLabels();
    ax = h->GetYaxis();
    ax->SetTitle("#epsilon_{#pi} [%]");
    ax->CenterTitle();
    ax->SetRangeUser(1.e-2, 10.);
    h->Draw();
    content = (TList *)fGraph->FindObject("Efficiencies");
    if(!(g = (TGraphErrors*)content->At(AliTRDpidUtil::kLQ))) break;
    if(!g->GetN()) break;
    legEff->SetHeader("PID Source");
    g->Draw("pc"); legEff->AddEntry(g, "LQ 2D", "pl");
    if(! (g = (TGraphErrors*)content->At(AliTRDpidUtil::kNN))) break;
    g->Draw("pc"); legEff->AddEntry(g, "NN", "pl");
    if(! (g = (TGraphErrors*)content->At(AliTRDpidUtil::kESD))) break;
    g->Draw("p"); legEff->AddEntry(g, "ESD", "pl");
    legEff->Draw();
    gPad->SetLogy();
    gPad->SetLogx();
    gPad->SetGridy();
    gPad->SetGridx();

    pad = ((TVirtualPad*)l->At(1));pad->cd();
    pad->SetMargin(0.1, 0.01, 0.1, 0.01);
    h=new TH1S("hThr", "", 1, .5, 11.);
    h->SetLineColor(1);h->SetLineWidth(1);
    ax = h->GetXaxis();
    ax->SetTitle("p [GeV/c]");
    ax->SetMoreLogLabels();
    ax = h->GetYaxis();
    ax->SetTitle("Threshold [%]");
    ax->SetRangeUser(5.e-2, 1.);
    h->Draw();
    content = (TList *)fGraph->FindObject("Thresholds");
    if(!(g = (TGraphErrors*)content->At(AliTRDpidUtil::kLQ))) break;
    if(!g->GetN()) break;
    g->Draw("pc");
    if(!(g = (TGraphErrors*)content->At(AliTRDpidUtil::kNN))) break;
    g->Draw("pc");
    if(!(g = (TGraphErrors*)content->At(AliTRDpidUtil::kESD))) break;
    g->Draw("p");
    gPad->SetLogx();
    gPad->SetGridy();
    gPad->SetGridx();
    return kTRUE;
  }
  case kdEdx:{
    // save 2.0 GeV projection as reference
    TLegend *legdEdx = new TLegend(.7, .7, .98, .98);
    legdEdx->SetBorderSize(1);
    kFIRST = kTRUE;
    if(!(h2 = (TH2F*)(fContainer->At(kdEdx)))) break;
    legdEdx->SetHeader("Particle Species");
    gPad->SetMargin(0.1, 0.01, 0.1, 0.01);
    for(Int_t is = AliPID::kSPECIES-1; is>=0; is--){
      Int_t bin = FindBin(is, 2.);
      h1 = h2->ProjectionY(Form("px%d", is), bin, bin);
      if(!h1->GetEntries()) continue;
      h1->Scale(1./h1->Integral());
      h1->SetLineColor(AliTRDCalPID::GetPartColor(is));
      if(kFIRST){
        h1->GetXaxis()->SetTitle("dE/dx (a.u.)");
        h1->GetYaxis()->SetTitle("<Entries>");
      }
      h = (TH1F*)h1->DrawClone(kFIRST ? "c" : "samec");
      legdEdx->AddEntry(h, Form("%s", AliTRDCalPID::GetPartName(is)), "l");
      kFIRST = kFALSE;
    }
    if(kFIRST) break;
    legdEdx->Draw();
    gPad->SetLogy();
    gPad->SetLogx(0);
    gPad->SetGridy();
    gPad->SetGridx();
    return kTRUE;
  }
  case kdEdxSlice:
    break;
  case kPH:{
    gPad->Divide(2, 1, 1.e-5, 1.e-5);
    TList *l=gPad->GetListOfPrimitives();

    // save 2.0 GeV projection as reference
    TLegend *legPH = new TLegend(.4, .7, .68, .98);
    legPH->SetBorderSize(1);legPH->SetFillColor(0);
    legPH->SetHeader("Particle Species");
    if(!(arr = (TObjArray*)(fContainer->At(kPH)))) break;
    if(!(h2 = (TProfile2D*)(arr->At(0)))) break;

    TVirtualPad *pad = ((TVirtualPad*)l->At(0));pad->cd();
    pad->SetMargin(0.1, 0.01, 0.1, 0.01);
    kFIRST = kTRUE;
    for(Int_t is=0; is<AliPID::kSPECIES; is++){
      Int_t bin = FindBin(is, 2.);
      h1 = h2->ProjectionY(Form("pyt%d", is), bin, bin);
      if(!h1->GetEntries()) continue;
      h1->SetMarkerStyle(24);
      h1->SetMarkerColor(AliTRDCalPID::GetPartColor(is));
      h1->SetLineColor(AliTRDCalPID::GetPartColor(is));
      if(kFIRST){
        h1->GetXaxis()->SetTitle("t_{drift} [100*ns]");
        h1->GetYaxis()->SetTitle("<dQ/dt> [a.u.]");
      }
      h = (TH1F*)h1->DrawClone(kFIRST ? "c" : "samec");
      legPH->AddEntry(h, Form("%s", AliTRDCalPID::GetPartName(is)), "pl");
      kFIRST = kFALSE;
    }

    pad = ((TVirtualPad*)l->At(1));pad->cd();
    pad->SetMargin(0.1, 0.01, 0.1, 0.01);
    if(!(h2 = (TProfile2D*)(arr->At(1)))) break;
    kFIRST = kTRUE;
    for(Int_t is=0; is<AliPID::kSPECIES; is++){
      Int_t bin = FindBin(is, 2.);
      h1 = h2->ProjectionY(Form("pyx%d", is), bin, bin);
      if(!h1->GetEntries()) continue;
      h1->SetMarkerStyle(24);
      h1->SetMarkerColor(AliTRDCalPID::GetPartColor(is));
      h1->SetLineColor(AliTRDCalPID::GetPartColor(is));
      if(kFIRST){
        h1->GetXaxis()->SetTitle("x_{drift} [cm]");
        h1->GetYaxis()->SetTitle("<dQ/dl> [a.u./cm]");
      }
      h = (TH1F*)h1->DrawClone(kFIRST ? "c" : "samec");
      kFIRST = kFALSE;
    }

    if(kFIRST) break;
    legPH->Draw();
    gPad->SetLogy(0);
    gPad->SetLogx(0);
    gPad->SetGridy();
    gPad->SetGridx();
    return kTRUE;
  }
  case kNClus:{
    // save 2.0 GeV projection as reference
    TLegend *legNClus = new TLegend(.13, .7, .4, .98);
    legNClus->SetBorderSize(1);
    legNClus->SetFillColor(0);

    kFIRST = kTRUE;
    if(!(h2 = (TH2F*)(fContainer->At(kNClus)))) break;
    legNClus->SetHeader("Particle Species");
    for(Int_t is=0; is<AliPID::kSPECIES; is++){
      Int_t bin = FindBin(is, 2.);
      h1 = h2->ProjectionY(Form("pyNClus%d", is), bin, bin);
      if(!h1->GetEntries()) continue;
      h1->Scale(100./h1->Integral());
      //h1->SetMarkerStyle(24);
      //h1->SetMarkerColor(AliTRDCalPID::GetPartColor(is));
      h1->SetLineColor(AliTRDCalPID::GetPartColor(is));
      if(kFIRST){ 
        h1->GetXaxis()->SetTitle("N^{cl}/tracklet");
        h1->GetYaxis()->SetTitle("Prob. [%]");
        h = (TH1F*)h1->DrawClone("c");
        h->SetMaximum(55.);
        h->GetXaxis()->SetRangeUser(0., 35.);
        kFIRST = kFALSE;
      } else h = (TH1F*)h1->DrawClone("samec");

      legNClus->AddEntry(h, Form("%s", AliTRDCalPID::GetPartName(is)), "l");
    }
    if(kFIRST) break;
    legNClus->Draw();
    gPad->SetLogy(0);
    gPad->SetLogx(0);
    gPad->SetGridy();
    gPad->SetGridx();
    return kTRUE;
  }
  case kMomentum:
  case kMomentumBin:
    break; 
  case kNTracklets:{
    TLegend *legNClus = new TLegend(.4, .7, .68, .98);
    legNClus->SetBorderSize(1);
    kFIRST = kTRUE;
    if(!(h2 = (TH2F*)(fContainer->At(kNTracklets)))) break;
    legNClus->SetHeader("Particle Species");
    for(Int_t is=0; is<AliPID::kSPECIES; is++){
      Int_t bin = FindBin(is, 2.);
      h1 = h2->ProjectionY(Form("pyNTracklets%d", is), bin, bin);
      if(!h1->GetEntries()) continue;
      h1->Scale(100./h1->Integral());
      //h1->SetMarkerStyle(24);
      //h1->SetMarkerColor(AliTRDCalPID::GetPartColor(is));
      h1->SetLineColor(AliTRDCalPID::GetPartColor(is));
      if(kFIRST){ 
        h1->GetXaxis()->SetTitle("N^{trklt}/track");
        h1->GetXaxis()->SetRangeUser(1.,6.);
        h1->GetYaxis()->SetTitle("Prob. [%]");
      }
      h = (TH1F*)h1->DrawClone(kFIRST ? "c" : "samec");
      legNClus->AddEntry(h, Form("%s", AliTRDCalPID::GetPartName(is)), "l");
      kFIRST = kFALSE;
    }
    if(kFIRST) break;
    legNClus->Draw();
    gPad->SetLogy(0);
    gPad->SetLogx(0);
    gPad->SetGridy();
    gPad->SetGridx();
    return kTRUE;
  }
  }
  AliInfo(Form("Reference plot [%d] missing result", ifig));
  return kFALSE;
}

//________________________________________________________________________
Bool_t AliTRDcheckPID::PostProcess()
{
  // Draw result to the screen
  // Called once at the end of the query

  if (!fContainer) {
    Printf("ERROR: list not available");
    return kFALSE;
  }
  if(!(fEfficiency = dynamic_cast<TObjArray *>(fContainer->At(kEfficiency)))){
    AliError("Efficiency container missing.");
    return 0x0;
  }
  if(!fGraph){ 
    fGraph = new TObjArray(6);
    fGraph->SetOwner();
    EvaluatePionEfficiency(fEfficiency, fGraph, 0.9);
  }
  fNRefFigures = 9;
  return kTRUE;
}

//________________________________________________________________________
void AliTRDcheckPID::EvaluatePionEfficiency(TObjArray * const histoContainer, TObjArray *results, Float_t electronEfficiency){
// Process PID information for pion efficiency

  fUtil->SetElectronEfficiency(electronEfficiency);

  Color_t colors[3] = {kBlue, kGreen+2, kRed};
  Int_t markerStyle[3] = {7, 7, 24};
  TString methodName[3] = {"LQ", "NN", "ESD"};
  // efficiency graphs
  TGraphErrors *g, *gPtrEff[3], *gPtrThres[3];
  TObjArray *eff = new TObjArray(3); eff->SetOwner(); eff->SetName("Efficiencies");
  results->AddAt(eff, 0);
  for(Int_t iMethod = 0; iMethod < 3; iMethod++){
    eff->AddAt(g = gPtrEff[iMethod] = new TGraphErrors(), iMethod);
    g->SetName(Form("efficiency_%s", methodName[iMethod].Data()));
    g->SetLineColor(colors[iMethod]);
    g->SetMarkerColor(colors[iMethod]);
    g->SetMarkerStyle(markerStyle[iMethod]);
  }

  // Threshold graphs
  TObjArray *thres = new TObjArray(3); thres->SetOwner(); thres->SetName("Thresholds");
  results->AddAt(thres, 1);
  for(Int_t iMethod = 0; iMethod < 3; iMethod++){
    thres->AddAt(g = gPtrThres[iMethod] = new TGraphErrors(), iMethod);
    g->SetName(Form("threshold_%s", methodName[iMethod].Data()));
    g->SetLineColor(colors[iMethod]);
    g->SetMarkerColor(colors[iMethod]);
    g->SetMarkerStyle(markerStyle[iMethod]);
  }
  
  Float_t mom = 0.;
  TH1D *histo1=0x0, *histo2=0x0;

  TH2F *hPtr[3];
  hPtr[0] = (TH2F*)histoContainer->At(AliTRDpidUtil::kLQ);
  hPtr[1] = (TH2F*)histoContainer->At(AliTRDpidUtil::kNN);
  hPtr[2] = (TH2F*)histoContainer->At(AliTRDpidUtil::kESD);
  
  for(Int_t iMom = 0; iMom < fMomentumAxis->GetNbins(); iMom++){
    mom = fMomentumAxis->GetBinCenter(iMom+1);

    Int_t binEl = fMomentumAxis->GetNbins() * AliPID::kElectron + iMom + 1, 
	  binPi = fMomentumAxis->GetNbins() * AliPID::kPion + iMom + 1;
    for(Int_t iMethod = 0; iMethod < 3; iMethod++){
      // Calculate the Pion Efficiency at 90% electron efficiency for each Method
      histo1 = hPtr[iMethod] -> ProjectionY(Form("%s_ele", methodName[iMethod].Data()), binEl, binEl);
      histo2 = hPtr[iMethod] -> ProjectionY(Form("%s_pio", methodName[iMethod].Data()), binPi, binPi);

      if(!fUtil->CalculatePionEffi(histo1, histo2)) continue;
     
      gPtrEff[iMethod]->SetPoint(iMom, mom, 1.e2*fUtil->GetPionEfficiency());
      gPtrEff[iMethod]->SetPointError(iMom, 0., 1.e2*fUtil->GetError());
      gPtrThres[iMethod]->SetPoint(iMom, mom, fUtil->GetThreshold());
      gPtrThres[iMethod]->SetPointError(iMom, 0., 0.);

      AliDebug(2, Form("Pion Efficiency for %s is : %f +/- %f", methodName[iMethod].Data(), fUtil->GetPionEfficiency(), fUtil->GetError()));
    }
  }
}