//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Tracks particles in given magnetic field
//
// Environment:
//      Software developed for the Prototype Detector at FOPI
//
// Author List:
//      Paul Buehler								<paul.buehler@oeaw.ac.at>
//
//-----------------------------------------------------------



// This Class' Header ------------------
#include "GEANEtrackerTask.h"

// C++ headers

// Collaborating Class Headers --------
#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "FOPIField.h"

#include "RKTrackRep.h"
#include "GFTrack.h"
#include "GFFieldManager.h"
#include "PndFieldAdaptor.h"
#include "GFException.h"

#include "TH2.h"
#include "TCanvas.h"
#include "TGraph.h"
#include "TMultiGraph.h"
#include "TLatex.h"
#include "TLegend.h"

#define DEBUG 0


//-----------------------------------------------------------
GEANEtrackerTask::GEANEtrackerTask() { }

//-----------------------------------------------------------
GEANEtrackerTask::~GEANEtrackerTask(){ }

//-----------------------------------------------------------
InitStatus GEANEtrackerTask::Init() {

  fprintf(stderr,"<I - GEANEtrackerTask::Init> Starting ...\n");

  //Get ROOT Manager ------------------
  FairRootManager* ioman= FairRootManager::Instance();
  if(ioman==0)
    {
      Error("GEANEtrackerTask::Init","RootManager not instantiated!");
      return kERROR;
    }
  
  // initialize propagator
  fPro = new FairGeanePro();
	
	// initialisation of initial conditions
	forig_err = TVector3(0., 0., 0.);
	fmom0_err = TVector3(0., 0., 0.);
	
 return kSUCCESS;
}

//-----------------------------------------------------------
void GEANEtrackerTask::SetParContainers() {

  fprintf(stderr,"<I - GEANEtrackerTask::SetParContainers> Starting ...\n");

  //  std::cout<<"PndTpcResidualTask::SetParContainers"<<std::endl;
  std::cout.flush();

  // Get run and runtime database
  FairRun* run = FairRun::Instance();
  if ( ! run ) Fatal("SetParContainers", "No analysis run");

  FairRuntimeDb* db = run->GetRuntimeDb();
  if ( ! db ) Fatal("SetParContainers", "No runtime database");

}

//-----------------------------------------------------------
void GEANEtrackerTask::Exec(Option_t* opt) {
	
  fprintf(stderr,"<I - GEANEtrackerTask::Exec> Starting ...\n");
  fprintf(stderr,"fmom0: %f %f %f\n",fmom0.X(),fmom0.Y(),fmom0.Z());

  // line colors
  char ltxt[200];
	Int_t lc[12] = {4,6,8,9,28,29,33,38,40,42,46,49};

	TVector3 mom = fmom0;
	TVector3 v1e = TVector3(1.,0.,0.);	
  TVector3 v2e = TVector3(0.,1.,0.);	
	TVector3 o, dj, dk;
	Double_t zfirst;
	TVector3 nextpos;
	Bool_t rc;
  
  RKTrackRep* rep = NULL;
  GFTrack*    trk = NULL;
  GFDetPlane plane;
  TVector3 GFpos, GFmom;
  TMatrixT<double> cov;  
  
	Char_t *pname = (Char_t*)fParticle->GetName();
	Int_t pdg = fParticle->PdgCode();
	Int_t pch = fParticle->Charge();
	
	// define start position
	Double_t vx = forig.X();
	Double_t vy = forig.Y();
	Double_t vz = forig.Z();
	
	// prepare histograms for plots
	Double_t const xmin =-150, xmax =  150;
	Double_t const ymin =-150, ymax =  150;
	Double_t const zmin =  -50, zmax =  500;
  TH2F hyx = TH2F("hyx","y versus x",10,ymin,ymax,10,xmin,xmax);
  hyx.SetStats(kFALSE);
  TH2F hzx = TH2F("hzx","z versus x",10,zmin,zmax,10,xmin,xmax);
  hzx.SetStats(kFALSE);
  TH2F hzy = TH2F("hzy","z versus y",10,zmin,zmax,10,ymin,ymax);
  hzy.SetStats(kFALSE);
	
  // Define z-range and number of steps for tracking
	if (forig.Z()<fzmin) {
	  zfirst = fzmin;
	} else {
	  zfirst = forig.Z();
	}
  Double_t dz = fzmax-zfirst;
	if (dz < 0) return;
	Int_t nsteps = (Int_t) dz/fzstep;
	if (nsteps < 1) return;
	
  // prepare TGraphs for all momenta to track
	TGraph *grxy = new TGraph[fnmomenta];
	TGraph *grzx = new TGraph[fnmomenta];
	TGraph *grzy = new TGraph[fnmomenta];
	TGraph *GFgrxy = new TGraph[fnmomenta];
	TGraph *GFgrzx = new TGraph[fnmomenta];
	TGraph *GFgrzy = new TGraph[fnmomenta];
	
	// prepare buffer for track points
	Double_t *xp = new Double_t[nsteps];
	Double_t *yp = new Double_t[nsteps];
	Double_t *zp = new Double_t[nsteps];
	Double_t *GFxp = new Double_t[nsteps];
	Double_t *GFyp = new Double_t[nsteps];
	Double_t *GFzp = new Double_t[nsteps];
	
	// prepare track parameters
	FairTrackParP *fStart = new FairTrackParP();
  FairTrackParP *fnow   =	new FairTrackParP();
  FairTrackParP *fRes   =	new FairTrackParP();
	
  // Initialise magnetic field
  FairRunAna* run = FairRunAna::Instance();
  FOPIField *fMagField = (FOPIField*) run->GetField();
  GFFieldManager::getInstance()->init(new PndFieldAdaptor(fMagField));

  // loop over momenta
	for (Int_t ii=0; ii<fnmomenta; ii++) {
    
		mom.SetMag(fmomenta[ii]);
		
    // Initialise GEANE
    *fStart = FairTrackParP(forig, mom, forig_err, fmom0_err, pch,forig,v1e,v2e);

    if (forig.Z()<fzmin) {
		  // get the starting point of the tracks at fzmin
		  nextpos = TVector3(0.,0.,fzmin);
	    fPro->PropagateToPlane(nextpos, v1e, v2e);
      rc =  fPro->Propagate(fStart, fnow, pdg);
	  } else {
	    fnow = fStart;
	  }
		Double_t zpos = fnow->GetZ()+fzstep;
    
    // Initialise GENFIT
    rep = new RKTrackRep(forig,mom,forig_err,fmom0_err,pdg);
    trk = new GFTrack(rep);
	  
		// loop over z-positions
		for (Int_t jj=0; jj<nsteps; jj++) {
    // fprintf(stderr,"xyz_in:  %f %f %f\n",fnow->GetX(),fnow->GetY(),fnow->GetZ());
	  
		  // GEANE propagation
      nextpos = TVector3(0.,0.,zpos);
	    fPro->PropagateToPlane(nextpos, v1e, v2e);
      rc =  fPro->Propagate(fnow, fRes, pdg);
		  fnow = fRes;
      // fprintf(stderr,"GEANE : Mom @ z=%f: %f\n",zpos,fRes->GetMomentum().Mag());
		  xp[jj] = fRes->GetX();
		  yp[jj] = fRes->GetY();
		  zp[jj] = fRes->GetZ();
		
		  
      // GENFIT propagation
      plane = GFDetPlane(
        nextpos,
        TVector3(1.,0.,0.),
        TVector3(0.,1.,0.));
      try {
        trk->getPosMomCov(plane,GFpos,GFmom,cov);
      } catch (GFException& e) {
        e.what();
      }
      // fprintf(stderr,"GENFIT: Mom @ z=%f: %f\n",zpos,GFmom.Mag());
		  GFxp[jj] = GFpos.X();
		  GFyp[jj] = GFpos.Y();
		  GFzp[jj] = GFpos.Z();
      
      zpos = zpos+fzstep;
	  }
	  grxy[ii] = TGraph(nsteps,xp,yp);
    grxy[ii].SetLineStyle(1);
    grxy[ii].SetLineWidth(1);
    grxy[ii].SetLineColor(lc[ii]);
	  grzx[ii] = TGraph(nsteps,zp,xp);
    grzx[ii].SetLineStyle(1);
    grzx[ii].SetLineWidth(1);
    grzx[ii].SetLineColor(lc[ii]);
	  grzy[ii] = TGraph(nsteps,zp,yp);
    grzy[ii].SetLineStyle(1);
    grzy[ii].SetLineWidth(1);
    grzy[ii].SetLineColor(lc[ii]);
	  
    GFgrxy[ii] = TGraph(nsteps,GFxp,GFyp);
    GFgrxy[ii].SetLineStyle(1);
    GFgrxy[ii].SetLineWidth(1);
    GFgrxy[ii].SetLineColor(1);
	  GFgrzx[ii] = TGraph(nsteps,GFzp,GFxp);
    GFgrzx[ii].SetLineStyle(1);
    GFgrzx[ii].SetLineWidth(1);
    GFgrzx[ii].SetLineColor(1);
	  GFgrzy[ii] = TGraph(nsteps,GFzp,GFyp);
    GFgrzy[ii].SetLineStyle(1);
    GFgrzy[ii].SetLineWidth(1);
    GFgrzy[ii].SetLineColor(1);
		
    fprintf(stderr,"xyz_out: %f %f %f\n",fRes->GetX(),fRes->GetY(),fRes->GetZ());
	}
	
	// plot track
  TCanvas * c1 = new TCanvas();
  TPad *pad1 = new TPad("pad1","1",0.05,0.52,0.48,0.97);
  TPad *pad2 = new TPad("pad2","2",0.52,0.52,0.97,0.97);
  TPad *pad3 = new TPad("pad3","3",0.05,0.05,0.48,0.48);
  TPad *pad4 = new TPad("pad4","4",0.52,0.05,0.97,0.48);
  pad1->Draw();
  pad2->Draw();
  pad3->Draw();
  pad4->Draw();
	
	// y - x
	TGraph *vgrxy = new TGraph(1,&vx,&vy);
  vgrxy->SetMarkerStyle(20);
  vgrxy->SetMarkerSize(0.3);
  vgrxy->SetMarkerColor(2);

  pad1->cd();
  hyx.Draw();
  TMultiGraph *mgrxy = new TMultiGraph();
  for (Int_t ii=0; ii<fnmomenta; ii++) {
    mgrxy->Add(&(grxy[ii]),"L");
    mgrxy->Add(&(GFgrxy[ii]),"L");
  }
  mgrxy->Add(vgrxy,"P");
  mgrxy->Draw("P");
	
  TAxis* tx1 = hyx.GetXaxis();
  tx1->SetTitle("x [cm]");
  tx1->SetNoExponent(kTRUE);
  TAxis* ty1 = hyx.GetYaxis();
  ty1->SetTitle("y [cm]");
  ty1->SetNoExponent(kTRUE);

	// z - y
	TGraph *vgrzy = new TGraph(1,&vz,&vy);
  vgrzy->SetMarkerStyle(20);
  vgrzy->SetMarkerSize(0.3);
  vgrzy->SetMarkerColor(2);

  pad2->cd();
  hzy.Draw();
  TMultiGraph *mgrzy = new TMultiGraph();
  for (Int_t ii=0; ii<fnmomenta; ii++) {
    mgrzy->Add(&(grzy[ii]),"L");
    mgrzy->Add(&(GFgrzy[ii]),"L");
  }
  mgrzy->Add(vgrzy,"P");
  mgrzy->Draw("L");
	
  TAxis* tx3 = hzy.GetXaxis();
  tx3->SetTitle("z [cm]");
  tx3->SetNoExponent(kTRUE);
  TAxis* ty3 = hzy.GetYaxis();
  ty3->SetTitle("y [cm]");
  ty3->SetNoExponent(kTRUE);

	// z - x
	TGraph *vgrzx = new TGraph(1,&vz,&vx);
  vgrzx->SetMarkerStyle(20);
  vgrzx->SetMarkerSize(0.3);
  vgrzx->SetMarkerColor(2);

  pad4->cd();
  hzx.Draw();
  TMultiGraph *mgrzx = new TMultiGraph();
  for (Int_t ii=0; ii<fnmomenta; ii++) {
    mgrzx->Add(&(grzx[ii]),"L");
    mgrzx->Add(&(GFgrzx[ii]),"L");
  }
  mgrzx->Add(vgrzx,"P");
  mgrzx->Draw("L");
	
  TAxis* tx2 = hzx.GetXaxis();
  tx2->SetTitle("z [cm]");
  tx2->SetNoExponent(kTRUE);
  TAxis* ty2 = hzx.GetYaxis();
  ty2->SetTitle("x [cm]");
  ty2->SetNoExponent(kTRUE);


  // legend
	pad3->cd();
  TLatex *txt = new TLatex();
  txt->SetTextSize(0.07);
  txt->SetNDC();
  
	Double_t theta = fmom0.Theta()/rad;
	Double_t phi   = fmom0.Phi()/rad;
	sprintf(ltxt, "%s #Theta / #phi [deg]: %7.2f / %7.2f",pname,theta, phi);
  txt->DrawLatex(0.05,0.9,ltxt);

	TLegend *leg = new TLegend(0.05,0.1,0.95,0.8);
  leg->AddEntry(vgrxy,"Origin of particle","P");
  for (Int_t ii=0; ii<fnmomenta; ii++) {
	  sprintf(ltxt,"GEANE %7.3f GeV/c", fmomenta[ii]);
		leg->AddEntry(&(grxy[ii]),ltxt,"L");
	  sprintf(ltxt,"GENFIT %7.3f GeV/c", fmomenta[ii]);
		leg->AddEntry(&(GFgrxy[ii]),ltxt,"L");
	}
  leg->Draw();

  // save canvas
  c1->Update();
  c1->Print("GEANEtracker_output.pdf");

	// cleanup
	c1->~TCanvas();
	delete xp, yp, zp;
}

ClassImp(GEANEtrackerTask)

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