// $Id$
// Main authors: Matevz Tadel & Alja Mrak-Tadel: 2006, 2007

/**************************************************************************
 * Copyright(c) 1998-2008, ALICE Experiment at CERN, all rights reserved. *
 * See http://aliceinfo.cern.ch/Offline/AliRoot/License.html for          *
 * full copyright notice.                                                 *
 **************************************************************************/

TEveTrack* esd_make_track(TEveTrackPropagator*   rnrStyle,
			  Int_t                  index,
			  AliESDtrack*           at,
			  AliExternalTrackParam* tp=0)
{
  // Helper function
  Double_t      pbuf[3], vbuf[3];
  TEveRecTrack  rt;

  if(tp == 0) tp = at;

  rt.fLabel  = at->GetLabel();
  rt.fIndex  = index;
  rt.fStatus = (Int_t) at->GetStatus();
  rt.fSign   = tp->GetSign();
  tp->GetXYZ(vbuf);
  rt.fV.Set(vbuf);
  tp->GetPxPyPz(pbuf);
  rt.fP.Set(pbuf);
  Double_t ep = at->GetP(), mc = at->GetMass();
  rt.fBeta = ep/TMath::Sqrt(ep*ep + mc*mc);

  TEveTrack* track = new TEveTrack(&rt, rnrStyle);
  //PH The line below is replaced waiting for a fix in Root
  //PH which permits to use variable siza arguments in CINT
  //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
  //PH  track->SetName(Form("ESDTrack %d", rt.fLabel));
  //PH  track->SetTitle(Form("pT=%.3f, pZ=%.3f; V=(%.3f, %.3f, %.3f)",
  //PH		       rt.fSign*TMath::Hypot(rt.fP.fX, rt.fP.fY), rt.fP.fZ,
  //PH		       rt.fV.fX, rt.fV.fY, rt.fV.fZ));
  char form[1000];
  sprintf(form,"TEveTrack %d", rt.fIndex);
  track->SetName(form);
  track->SetStdTitle();
  return track;
}

void esd_track_add_param(TEveTrack* track, AliExternalTrackParam* tp)
{
  if (tp == 0)
    return;

  Double_t      pbuf[3], vbuf[3];
  tp->GetXYZ(vbuf);
  tp->GetPxPyPz(pbuf);

  TEvePathMark pm(TEvePathMark::kReference);
  pm.fV.Set(vbuf);
  pm.fP.Set(pbuf);
  track->AddPathMark(pm);
}

// Use inner-tpc track params when its refit failed.
Bool_t gkFixFailedITSExtr    = kFALSE;

// Also show lines as generated by AliESDtrack.
Bool_t gkMakeTrackParamLines = kFALSE;

TEveTrackList* esd_tracks(Double_t min_pt=0.1, Double_t max_pt=100)
{
  AliESDEvent* esd = AliEveEventManager::AssertESD();

  Double_t minptsq = min_pt*min_pt;
  Double_t maxptsq = max_pt*max_pt;
  Double_t ptsq;

  TEveTrackList* cont = new TEveTrackList("ESD Tracks");
  cont->SetMainColor(Color_t(6));
  TEveTrackPropagator* rnrStyle = cont->GetPropagator();
  rnrStyle->SetMagField( 0.1*esd->GetMagneticField() );

  gEve->AddElement(cont);

  TEveElementList* contLines = 0;
  if (gkMakeTrackParamLines)
  {
    contLines = new TEveElementList("MyTracks");  
    gEve->AddElement(contLines);
  }

  Int_t    count = 0;
  Double_t pbuf[3];
  for (Int_t n=0; n<esd->GetNumberOfTracks(); n++)
  {
    AliESDtrack* at = esd->GetTrack(n);

    // Here would be sweet to have TObjectFormula.
    at->GetPxPyPz(pbuf);
    ptsq = pbuf[0]*pbuf[0] + pbuf[1]*pbuf[1];
    if(ptsq < minptsq || ptsq > maxptsq)
      continue;

    ++count;

    // If gkFixFailedITSExtr is TRUE (FALSE by default) and
    // if ITS refit failed, take track parameters at inner TPC radius.
    Bool_t innerTaken = kFALSE;
    AliExternalTrackParam* tp = at;
    if (gkFixFailedITSExtr && !at->IsOn(AliESDtrack::kITSrefit)) {
      tp = at->GetInnerParam();
      innerTaken = kTRUE;
    }

    TEveTrack* track = esd_make_track(rnrStyle, n, at, tp);
    track->SetAttLineAttMarker(cont);

    if (!innerTaken) {
      esd_track_add_param(track, at->GetInnerParam());
    }
    // esd_track_add_param(track, at->GetOuterParam());

    gEve->AddElement(track, cont);

    if (gkMakeTrackParamLines) {
      TEveLine* l = new TEveLine; 
      l->SetName(Form("Track%d", count));
      l->SetLineColor((Color_t)5);
      at->FillPolymarker(l, esd->GetMagneticField(), 0, 250, 5);
      contLines->AddElement(l);
    }
  }

  //PH The line below is replaced waiting for a fix in Root
  //PH which permits to use variable siza arguments in CINT
  //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
  //PH  const Text_t* tooltip = Form("pT ~ (%.2lf, %.2lf), N=%d", min_pt, max_pt, count);
  char tooltip[1000];
  sprintf(tooltip,"pT ~ (%.2lf, %.2lf), N=%d", min_pt, max_pt, count);
  cont->SetTitle(tooltip); // Not broadcasted automatically ...
  cont->UpdateItems();

  cont->MakeTracks();

  gEve->Redraw3D();

  return cont;
}

/******************************************************************************/
// esd_tracks_from_array()
/******************************************************************************/

TEveTrackList* esd_tracks_from_array(TCollection* col, AliESDEvent* esd=0)
{
  // Retrieves AliESDTrack's from collection.
  // See example usage with AliAnalysisTrackCuts in the next function.

  if (esd == 0) esd = AliEveEventManager::AssertESD();

  TEveTrackList* cont = new TEveTrackList("ESD Tracks");
  cont->SetMainColor(Color_t(6));
  TEveTrackPropagator* rnrStyle = cont->GetPropagator();
  rnrStyle->SetMagField( 0.1*esd->GetMagneticField() );

  gEve->AddElement(cont);

  Int_t    count = 0;
  TIter    next(col);
  TObject *obj;
  while ((obj = next()) != 0)
  {
    if (obj->IsA()->InheritsFrom("AliESDtrack") == kFALSE) {
      Warning("Object '%s', '%s' is not an AliESDtrack.",
	      obj->GetName(), obj->GetTitle());
      continue;
    }

    ++count;
    AliESDtrack* at = (AliESDtrack*) obj;

    TEveTrack* track = esd_make_track(rnrStyle, count, at);
    track->SetAttLineAttMarker(cont);
    gEve->AddElement(track, cont);
  }

  //PH The line below is replaced waiting for a fix in Root
  //PH which permits to use variable siza arguments in CINT
  //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
  //PH  const Text_t* tooltip = Form("N=%d", count);
  const tooltip[1000];
  sprintf(tooltip,"N=%d", count);
  cont->SetTitle(tooltip); // Not broadcasted automatically ...
  cont->UpdateItems();

  cont->MakeTracks();

  gEve->Redraw3D();

  return cont;
}

void esd_tracks_alianalcuts_demo()
{
  AliESDEvent* esd = AliEveEventManager::AssertESD();
  gSystem->Load("libANALYSIS");

  AliAnalysisTrackCuts atc;
  atc.SetPtRange(0.1, 5);
  atc.SetRapRange(-1, 1);

  esd_tracks_from_array(atc.GetAcceptedParticles(esd), esd);
}

/******************************************************************************/
// esd_tracks_vertex_cut
/******************************************************************************/

Float_t get_sigma_to_vertex(AliESDtrack* esdTrack)
{
  // Taken from: PWG0/esdTrackCuts/AliESDtrackCuts.cxx
  // Float_t AliESDtrackCuts::GetSigmaToVertex(AliESDtrack* esdTrack)

  Float_t b[2];
  Float_t bRes[2];
  Float_t bCov[3];
  esdTrack->GetImpactParameters(b,bCov);
  if (bCov[0]<=0 || bCov[2]<=0) {
    printf("Estimated b resolution lower or equal zero!\n");
    bCov[0]=0; bCov[2]=0;
  }
  bRes[0] = TMath::Sqrt(bCov[0]);
  bRes[1] = TMath::Sqrt(bCov[2]);

  // -----------------------------------
  // How to get to a n-sigma cut?
  //
  // The accumulated statistics from 0 to d is
  //
  // ->  Erf(d/Sqrt(2)) for a 1-dim gauss (d = n_sigma)
  // ->  1 - Exp(-d**2) for a 2-dim gauss (d*d = dx*dx + dy*dy != n_sigma)
  //
  // It means that for a 2-dim gauss: n_sigma(d) = Sqrt(2)*ErfInv(1 - Exp((-x**2)/2)
  // Can this be expressed in a different way?

  if (bRes[0] == 0 || bRes[1] ==0)
    return -1;

  Float_t d = TMath::Sqrt(TMath::Power(b[0]/bRes[0],2) + TMath::Power(b[1]/bRes[1],2));

  // stupid rounding problem screws up everything:
  // if d is too big, TMath::Exp(...) gets 0, and TMath::ErfInverse(1) that should be infinite, gets 0 :(
  if (TMath::Exp(-d * d / 2) < 1e-10)
    return 1000;

  d = TMath::ErfInverse(1 - TMath::Exp(-d * d / 2)) * TMath::Sqrt(2);
  return d;
}

TEveElementList* esd_tracks_vertex_cut()
{
  // Import ESD tracks, separate them into five containers according to
  // primary-vertex cut and ITS refit status.

  AliESDEvent* esd = AliEveEventManager::AssertESD();

  TEveElementList* cont = new TEveElementList("ESD Tracks", 0, kTRUE);
  gEve->AddElement(cont);
  TEveTrackList *tl[5];
  Int_t            tc[5];
  Int_t            count = 0;

  tl[0] = new TEveTrackList("Sigma < 3");
  tc[0] = 0;
  tl[0]->GetPropagator()->SetMagField( 0.1*esd->GetMagneticField() );
  tl[0]->SetMainColor(Color_t(3));
  gEve->AddElement(tl[0], cont);

  tl[1] = new TEveTrackList("3 < Sigma < 5");
  tc[1] = 0;
  tl[1]->GetPropagator()->SetMagField( 0.1*esd->GetMagneticField() );
  tl[1]->SetMainColor(Color_t(7));
  gEve->AddElement(tl[1], cont);

  tl[2] = new TEveTrackList("5 < Sigma");
  tc[2] = 0;
  tl[2]->GetPropagator()->SetMagField( 0.1*esd->GetMagneticField() );
  tl[2]->SetMainColor(Color_t(46));
  gEve->AddElement(tl[2], cont);

  tl[3] = new TEveTrackList("no ITS refit; Sigma < 5");
  tc[3] = 0;
  tl[3]->GetPropagator()->SetMagField( 0.1*esd->GetMagneticField() );
  tl[3]->SetMainColor(Color_t(41));
  gEve->AddElement(tl[3], cont);

  tl[4] = new TEveTrackList("no ITS refit; Sigma > 5");
  tc[4] = 0;
  tl[4]->GetPropagator()->SetMagField( 0.1*esd->GetMagneticField() );
  tl[4]->SetMainColor(Color_t(48));
  gEve->AddElement(tl[4], cont);

  for (Int_t n=0; n<esd->GetNumberOfTracks(); n++)
  {
    AliESDtrack* at = esd->GetTrack(n);

    Float_t s  = get_sigma_to_vertex(at);
    Int_t   ti;
    if      (s <  3) ti = 0;
    else if (s <= 5) ti = 1;
    else             ti = 2;

    AliExternalTrackParam* tp = at;
    // If ITS refit failed, optionally take track parameters at inner
    // TPC radius and put track in a special container.
    // This ignores state of gkFixFailedITSExtr (used in esd_tracks()).
    // Use BOTH functions to compare results.
    if (!at->IsOn(AliESDtrack::kITSrefit)) {
      // tp = at->GetInnerParam();
      ti = (ti == 2) ? 4 : 3;
    }

    TEveTrackList* tlist = tl[ti];
    ++tc[ti];
    ++count;

    TEveTrack* track = esd_make_track(tlist->GetPropagator(), n, at, tp);
    track->SetAttLineAttMarker(tlist);

    //PH The line below is replaced waiting for a fix in Root
    //PH which permits to use variable siza arguments in CINT
    //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
    //PH    track->SetName(Form("track %d, sigma=%5.3f", at->GetLabel(), s));
    char form[1000];
    sprintf(form,"TEveTrack idx=%d, sigma=%5.3f", at->GetID(), s);
    track->SetName(form);
    gEve->AddElement(track, tlist);
  }

  for (Int_t ti=0; ti<5; ++ti) {
    TEveTrackList* tlist = tl[ti];
    //PH The line below is replaced waiting for a fix in Root
    //PH which permits to use variable siza arguments in CINT
    //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
    //PH    const Text_t* tooltip = Form("N tracks=%d", tc[ti]);
    //MT Modified somewhat.
    char buff[1000];
    sprintf(buff, "%s [%d]", tlist->GetName(), tlist->GetNChildren());
    tlist->SetName(buff);
    sprintf(buff, "N tracks=%d", tc[ti]);
    tlist->SetTitle(buff); // Not broadcasted automatically ...
    tlist->UpdateItems();

    tlist->MakeTracks();
  }
  //PH The line below is replaced waiting for a fix in Root
  //PH which permits to use variable siza arguments in CINT
  //PH on some platforms (alphalinuxgcc, solariscc5, etc.)
  //PH  cont->SetTitle(Form("N all tracks = %d", count));
  char form[1000];
  sprintf(form,"N all tracks = %d", count);
  cont->SetTitle(form);
  cont->UpdateItems();
  gEve->Redraw3D();

  return cont;
}