//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class TpcGenfitTestTask
//      see TpcGenfitTestTask.hh for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Christian Hoeppner     TUM         
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

// This Class' Header ------------------
#include "TpcGenfitTestTask.h"
#include <cmath>
// C/C++ Headers ----------------------
#include <iostream>
#include <assert.h> 

// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "TClonesArray.h"
#include "GFTrack.h"
#include "TpcCluster.h"
#include "TpcDigiMapper.h"
#include "TpcFrontend.h"
#include "TpcSPHit.h"
#include "TpcTestPlanarHit.h"
#include "PndMCTrack.h"
#include "TpcPoint.h"


#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "TpcDigiPar.h"

#include "GeaneTrackRep.h"
#include "RKTrackRep.h"
#include "GFAbsTrackRep.h"
#include "GFRecoHitFactory.h"
#include "GFKalman.h"
#include "GFException.h"
#include "GFTrack.h"
#include "TFile.h"
#include "GFDetPlane.h"
#include "GFAbsRecoHit.h"
#include "GFFieldManager.h"

#include"PndFieldAdaptor.h"

#include "TVector3.h"
#include "TRandom3.h"
#include "TGeoManager.h"
#include "TCanvas.h"
#include "TH1D.h"
#include "TStyle.h"
#include "TDatabasePDG.h"

#include <signal.h>
#include <stdlib.h>

#define NREPS 2

// Class Member definitions -----------


TpcGenfitTestTask::TpcGenfitTestTask()
  : FairTask("GENFIT test Task"),
  minDist(1.)
{

}


TpcGenfitTestTask::~TpcGenfitTestTask()
{
}

InitStatus
TpcGenfitTestTask::Init()
{
  //Get ROOT Manager
  FairRootManager* ioman= FairRootManager::Instance();
  
  if(ioman==0) {
    Error("TpcGenfitTestTask::Init","RootManager not instantiated!");
    return kERROR;
  }
  
  // Get input collection
  _pointArray=(TClonesArray*) ioman->GetObject("TpcPoint");
  
  if(_pointArray==0) {
    Error("TpcGenfitTestTask::Init","Point-array not found!");
    return kERROR;
  }
    
  _mcTrackArray=(TClonesArray*) ioman->GetObject("MCTrack");
  if(_mcTrackArray==0) {
    Error("TpcGenfitTestTask::Init","MC-Track-array not found!");
    return kERROR;
  }


  int nBins(200);
  double width(10);

  for (unsigned int i=0; i<NREPS; ++i){
    hqopPu.push_back(new TH1D("hqopPu","q/p pull",nBins, -1*width, width));
    hupPu.push_back(new TH1D("hupPu","u' pull",nBins, -1*width, width));
    hvpPu.push_back(new TH1D("hvpPu","v' pull",nBins, -1*width, width));
    huPu.push_back(new TH1D("huPu","u pull",nBins, -1*width, width));
    hvPu.push_back(new TH1D("hvPu","v pull",nBins, -1*width, width));

    hChi2.push_back(new TH1D("hChi2","chi2",200,0,200));
    hRedChi2.push_back(new TH1D("hRedChi2","reduced chi2",50,0,5));
    hNDF.push_back(new TH1D("hNDF","NDF",200,0,200));
    hNFail.push_back(new TH1D("hNFail","n failed hits",20,0,20));
  }


  //get the magnetic field for curvature seeding
  FairField* field=FairRunAna::Instance()->GetField();
  GFFieldManager::getInstance()->init(new PndFieldAdaptor(field));
  _geanePro=new FairGeanePro();
  gGeoManager->Print();
  return kSUCCESS;
}


void
TpcGenfitTestTask::SetParContainers() {

  // 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
TpcGenfitTestTask::Exec(Option_t* opt)
{
  static int counter(0);
  static TRandom3 rand(0);

  static const TVector3 resolution(0.03,0.03,0.1); // cm
  static const TVector3 momresolution(0.1,0.1,0.1); // relative

  std::cout << "TpcGenfitTestTask::Exec " << counter++ << std::endl;

  
  //loop over MC points and collect 
  std::vector<TpcPoint*> pointlist;
  for(int p=0; p<_pointArray->GetEntriesFast(); ++p) {
    int id = ((TpcPoint*)_pointArray->At(p))->GetTrackID();
    //only primary tracks
    if(id==0) pointlist.push_back((TpcPoint*)_pointArray->At(p));
  }

  if(pointlist.size()<4.) return;

  // get positions with a mimimum distance
  std::vector<TVector3> positions;
  TVector3 vec;
  pointlist[0]->Position(vec);
  positions.push_back(vec);
  TVector3 before(vec);

  for(unsigned int i=1; i<pointlist.size(); ++i){
    pointlist[i]->Position(vec);
    double dist = (vec-before).Mag();
    if(dist>minDist){
      before=vec;
      positions.push_back(vec);
    }
  }

  if(positions.size()<4.) return;

  
  // get mc pdg code and charge
  int pdgid = ((PndMCTrack*)_mcTrackArray->At(0))->GetPdgCode();

  TParticlePDG * part = TDatabasePDG::Instance()->GetParticle(pdgid);
  if(part == 0) return;
  double charge = part->Charge()/(3.);


  // get momentum
  TVector3 mom;
  pointlist[0]->Momentum(mom);



  // reference plane & rep
  GFDetPlane plRef(positions[0],mom);

  GFAbsTrackRep *repRef = new RKTrackRep(positions[0],
                                         mom,
                                         TVector3(1.,1.,1.),
                                         TVector3(.2,.2,.2),
                                         pdgid);
  

  // smear seed values
  TVector3 momerr(momresolution.X()*mom.X(),
                  momresolution.Y()*mom.Y(),
                  momresolution.Z()*mom.Z());

  TVector3 pos(positions[0]);
  pos.SetXYZ(rand.Gaus(pos.X(), resolution.X()),
             rand.Gaus(pos.Y(), resolution.Y()),
             rand.Gaus(pos.Z(), resolution.Z()));
  mom.SetXYZ(rand.Gaus(mom.X(), momerr.X()),
             rand.Gaus(mom.Y(), momerr.Y()),
             rand.Gaus(mom.Z(), momerr.Z()) );
 

  // create trackreps
  std::vector<GFAbsTrackRep*> reps;

  reps.push_back(new RKTrackRep(pos,
                                mom,
                                resolution,
                                momerr,
                                pdgid));
  ((RKTrackRep*)(reps.back()))->disableMaterialEffects();

  reps.push_back(new GeaneTrackRep(_geanePro,
                                   GFDetPlane(pos, mom),
                                   mom,
                                   resolution,
                                   momerr,
                                   charge,
                                   pdgid));

  assert (reps.size() == NREPS);

  // create track and add reps
  GFTrack trk;

  for (unsigned int i=0; i<reps.size(); ++i){
    trk.addTrackRep(reps[i]);
  }

  // add hits (and do smearing)
  // first hit is a planar hit in the reference plane
  trk.addHit(new TpcTestPlanarHit(plRef, resolution.X()), 3, 0);
  // other hits are space hits
  for(unsigned int i=1; i<positions.size(); ++i){
    trk.addHit(new TpcSPHit(positions.at(i), resolution), 3, i);
  }


  // fitting
  static GFKalman k;
  try {
    k.processTrack(&trk);

    // make sure we are at the reference plane!
    for (unsigned int i=0; i<reps.size(); ++i){
      if (reps[i]->getReferencePlane() != plRef) reps[i]->extrapolate(plRef);
    }
  }
  catch(GFException & e) {
    std::cerr << "exception caught, skipping event ..." << std::endl;
    std::cout << e.what() << std::endl;
    delete repRef;
    return;
  }


  // calc pulls
  TMatrixT<double> referenceState(repRef->getState());
  delete repRef;

  for (unsigned int i=0; i<reps.size(); ++i){

    TMatrixT<double> state(reps[i]->getState());
    TMatrixT<double> cov(reps[i]->getCov());

    hqopPu[i]->Fill( (state[0][0]-referenceState[0][0]) / sqrt(cov[0][0]) );
    hupPu[i]->Fill(  (state[1][0]-referenceState[1][0]) / sqrt(cov[1][1]) );
    hvpPu[i]->Fill(  (state[2][0]-referenceState[2][0]) / sqrt(cov[2][2]) );
    huPu[i]->Fill(   (state[3][0]-referenceState[3][0]) / sqrt(cov[3][3]) );
    hvPu[i]->Fill(   (state[4][0]-referenceState[4][0]) / sqrt(cov[4][4]) );

    hChi2[i]->Fill(   reps[i]->getChiSqu() );
    hRedChi2[i]->Fill(reps[i]->getRedChiSqu());
    hNDF[i]->Fill(    reps[i]->getNDF()    );
    hNFail[i]->Fill(  trk.getFailedHits(i) );
  }

}


void 
TpcGenfitTestTask::WriteTree(TString outname){
  std::cout << "WriteTree" << std::endl;

  outfile = TFile::Open(outname, "RECREATE");

  gROOT->SetStyle("Plain");
  gStyle->SetPalette(1);
  gStyle->SetOptFit(1111);

  outfile->cd();

  // fit and draw histograms

  for (unsigned int i=0; i<NREPS; ++i){

    unsigned int j(0);
    TCanvas* c1 = new TCanvas();
    c1->Divide(3,3);

    c1->cd(++j);
    huPu[i]->Fit("gaus", "L M");
    huPu[i]->Draw();

    c1->cd(++j);
    hvPu[i]->Fit("gaus", "L M");
    hvPu[i]->Draw();

    c1->cd(++j);
    hqopPu[i]->Fit("gaus", "L M");
    hqopPu[i]->Draw();


    c1->cd(++j);
    hupPu[i]->Fit("gaus", "L M");
    hupPu[i]->Draw();

    c1->cd(++j);
    hvpPu[i]->Fit("gaus", "L M");
    hvpPu[i]->Draw();

    c1->cd(++j);
    hNFail[i]->Draw();


    c1->cd(++j);
    hChi2[i]->Draw();

    c1->cd(++j);
    hNDF[i]->Draw();

    c1->cd(++j);
    hRedChi2[i]->Draw();

    c1->Write();
  }

  outfile->Close();

 std::cout << "Pulls output file is "    << outname << std::endl;
}



ClassImp(TpcGenfitTestTask)