//-----------------------------------------------------------------------------
//  Dec 26 2000 P.Murat: initialization of the STNTUPLE track block
//-----------------------------------------------------------------------------
#include "TROOT.h"
#include "TFolder.h"

#include <AbsEnv/AbsEnv.hh>

#include "Electron/emobj_alg.hh"

#include <TrackingObjects/Storable/CdfTrackColl.hh>
#include <TrackingObjects/Storable/CdfTrackView.hh>
#include <TrackingObjects/Storable/CdfTrackMatch.hh>
#include "TrackingCT/CT_ReFit.hh"
#include <PhoenixMods/PhoenixTools.hh>
#include <Stntuple/mod/StntupleCotReFit.hh>
#include "Stntuple/mod/StntupleSiExpected.hh"
#include "TrackingSI/Utils/SiExpected.hh"
#include "TrackingHL/Utility/Cot_tZero.hh"
#include <Stntuple/mod/StntupleUtilities.hh>
#include <Stntuple/alg/TStntuple.hh>

#include <Stntuple/obj/TStnEvent.hh>
#include <Stntuple/obj/TStnErrorLogger.hh>
#include <Stntuple/obj/TStnTrackBlock.hh>
#include <Stntuple/obj/TStnVertexBlock.hh>
#include <Stntuple/mod/InitStntupleDataBlocks.hh>

#include "SimulationObjects/OBSP_StorableBank.hh"

#include "Muon/MuonFiducialTool.hh"
//_____________________________________________________________________________
Int_t StntupleInitTrack(TStnTrack* trk, const CdfTrack* rick, Int_t Mode) 
{
  // initialize STNTUPLE track starting from CDF track, should be called from
  // only StntupleInitTrackBlock

  static TFolder*         fol       = NULL;
  static TStnErrorLogger* logger    = NULL;
  static CT_ReFit*        refit     = NULL;
  static SiExpected*      siexpected= NULL;

  StntupleSiExpected*     stnsiexp  = NULL;
  StntupleCotReFit*       scrf      = NULL;
  const CdfTrack*         CotParent = NULL;

  if (! fol) {
//-----------------------------------------------------------------------------
//  initialize local static variables
//-----------------------------------------------------------------------------
    fol     = (TFolder*) gROOT->GetRootFolder()->FindObject("StntupleFolder");
    logger  = (TStnErrorLogger*) fol->FindObject("StntupleErrorLogger");
    scrf    = (StntupleCotReFit*) fol->FindObject("CotReFit");
    refit   = scrf->GetReFit();
    stnsiexp= (StntupleSiExpected*)  fol->FindObject("SiExpected");
    siexpected = stnsiexp->GetSiExpected();
  }

  HepLorentzVector   p(rick->fourMomentum(0.1396));

  trk->fAlgorithm = 0;
  
  trk->fAlgorithm = ( (rick->algorithm().value() ) |
		      (rick->algorithm().CotParentAlgo() << 8) );

  //Add Muon fiducial tool for track with pt>1.5
  if(rick->pt()>1.5){
    MuonFiducialTool fidTool;
    UChar_t CMUFid = 0, CMUFid2 = 0;
    UChar_t CMPFid = 0, CMPFid2 = 0;
    UChar_t CMXFid = 0, CMXFid2 = 0;
    UChar_t BMUFid = 0, BMUFid2 = 0;
    double d=0,x=0,z=0;
    if(fidTool.isFiducial(*rick,MuonDigiCode::CMU,d,x,z)) {
      CMUFid=1;
      if(x<0.0 && z<-3.0) CMUFid2=1;
    }
    if(fidTool.isFiducial(*rick,MuonDigiCode::CMP,d,x,z)) {
      CMPFid=1;
      if(x<0.0 && z<-3.0) CMPFid2=1;
    }
    if(fidTool.isFiducial(*rick,MuonDigiCode::CMX,d,x,z)) {
      CMXFid=1;
      if(x<0.0 && z<-3.0) CMXFid2=1;
    }
    if(fidTool.isFiducial(*rick,MuonDigiCode::BMU,d,x,z)) {
      BMUFid=1;
      if(x<0.0 && z<-3.0) BMUFid2=1;
    }   

    trk->fAlgorithm =( trk->fAlgorithm  |
		       (CMUFid << 16) |
		       (CMPFid << 17) |
		       (CMXFid << 18) |
		       (BMUFid << 19) );
    trk->fAlgorithm =( trk->fAlgorithm  |
		       (CMUFid2 << 20) |
		       (CMPFid2 << 21) |
		       (CMXFid2 << 22) |
		       (BMUFid2 << 23) );

    // Add fidEle(track) code (e.g. for stubless muon finding) :
    trk->fAlgorithm =( trk->fAlgorithm  |
		       (fidEle(*rick) << 24 & 0xff000000) );
  }

  trk->fMomentum.SetXYZT(p.px(),p.py(),p.pz(),p.e());

  trk->fVind    = -1;

  trk->fSvxHitWord  = ( (rick->numSIHits   ()      ) |
			(rick->numSIHitsPhi() <<  8) |
			(rick->numSIHitsSt () << 16) |
			(rick->numSIHitsZ  () << 24)    );
  trk->fSvxHitMask = rick->usedSI();

//-----------------------------------------------------------------------------
// trace back rick's ultimate Cot parent
//-----------------------------------------------------------------------------
  CotParent = rick;
  while(CotParent->parent().is_nonnull()) CotParent = &*(CotParent->parent());
  if(CotParent->usedSI()) CotParent = NULL; //need to do this when rick is Si only track
                                            //So we don't try to do a CT_Refit
//-----------------------------------------------------------------------------
// number of COT segments: define segment as >= 6 hits in the superlayer
//-----------------------------------------------------------------------------
  int       nax_seg, nst_seg;
  nst_seg = 0;
  nax_seg = 0;

  CT_LayerMask ricks_mask =  rick->usedCT();
  
  for(int i=0; i<3; i++){
    trk->fCotHitMask[i]=0;
    for(int j=0; j<32; j++){
      int k=i*32+j;
      int bit = ricks_mask[k];
      trk->fCotHitMask[i] |= (bit << j);
    }
  }
  
  for (int i=0; i<8; i+=2) {
    if (rick->numCTHits(i) >= 5) nst_seg++;
  }
					// number of axial SL's with hits
  for (int i=1; i<8; i+=2) {
    if (rick->numCTHits(i) >= 5) nax_seg++;
  }
					// byte for COT segments

  int seg  = ((nst_seg << 4) + nax_seg) & 0xff;

  trk->fCotHitWord  = ( ( rick->numCTHits  ()      ) | 
			( rick->numCTHitsAx() <<  8) | 
			( rick->numCTHitsSt() << 16) |
			( seg                 << 24)    );



  trk->fCot     = rick->lambda();	// cot(theta) of the track
  trk->fCurv    = rick->curvature();
  trk->fZ0      = rick->z0();		// z_0 of the track
  trk->fD0      = rick->d0();
  trk->fPhi0    = rick->phi0();
//-----------------------------------------------------------------------------
// T0
//-----------------------------------------------------------------------------
  trk->fT0      = -1.e6;
  trk->fT0Sigma =  1.e4;
  CdfTrackView* oneTrackView = new CdfTrackView();
  oneTrackView->contents().push_back(rick);
  Cot_tZero tZero;
  tZero.calculate_t0(oneTrackView->contents());
  if(tZero.track_t0().size() > 0){
    trk->fT0  = tZero.track_t0()[0];
    trk->fT0Sigma = tZero.sigma_track_t0()[0];
  }
  oneTrackView->destroy();

//-----------------------------------------------------------------------------
// track chi**2's: COT chi**2 is defined by the chi**2 of the parent
//-----------------------------------------------------------------------------
  trk->fChi2    = rick->chi2();

  if (CotParent) trk->fChi2Cot = CotParent->chi2();
  else           trk->fChi2Cot = TStnDataBlock::kUndefined;

  trk->fChi2Svx = rick->chi2SI();

  if (PhoenixTools::isPhoenixTrack(&(*rick))) {
    trk->fChi2Cot = PhoenixTools::chisqDofSeed(&(*rick));
  }

//-----------------------------------------------------------------------------
// track covariance matrix, do it quick and dirty first, then see how it can be
// improved (packed and such)
//-----------------------------------------------------------------------------
  for (int i=0; i<5; i++) {
    for (int j=0; j<5; j++) {
      trk->fCov(i,j) = rick->cov(i,j);
    }
  }

  trk->fPt      = rick->pt();
  trk->fEta     = trk->fMomentum.Eta();
//-----------------------------------------------------------------------------
// beam constrained parameters
//-----------------------------------------------------------------------------
  trk->fBcvind  = 0;
  trk->fBcZ0    = -1.e6;
  trk->fBcD0    = -1.e6;
  trk->fBcPhi0  = -1.e6;
  trk->fBcC0    = -1.e6;
  trk->fBcLam0  = -1.e6;
  
  for (int i=0; i<4; i++) trk->fBcp4[i] = 1.e-6;
//-----------------------------------------------------------------------------
// and tracking isolation
//-----------------------------------------------------------------------------
  trk->fIso4    = +1.e6;
//-----------------------------------------------------------------------------
// and Si hit isolation
//-----------------------------------------------------------------------------

  siexpected->setCountHitsCone( true );
  double conesize=0.4;
  siexpected->setCone( conesize );
  
  std::vector<SiExpected::Info> exphits;
  exphits.clear();
  int nSiHitsCone=0;
  int nSiHitsEx   =0;
  int nSiPhiHitsEx=0;
  if ( siexpected->findExpected ( *rick, exphits )) {
    for ( std::vector<SiExpected::Info>::iterator
	    expHitIter=exphits.begin(), expHitIterend = exphits.end () ;
	  expHitIter != expHitIterend ; ++expHitIter )
      {
        if ( (*expHitIter).integrated ) {
          nSiHitsEx++;
	  if ((*expHitIter).isPhi) nSiPhiHitsEx++;
	  nSiHitsCone+=(*expHitIter).nHitsCone;
        }
      }
  }
  if (nSiHitsCone>4096) {
    std::cout << "Warning: more than 4096 hits in 0.4 cone around track: nhit="<<nSiHitsCone<<std::endl;
    nSiHitsCone=4096;
  }
  trk->fSiExpWord=( (nSiHitsEx) | (nSiPhiHitsEx<<8) | (nSiHitsCone<<16) );

//-----------------------------------------------------------------------------
// Anyes: do beam-constrained fit of the track parameters. 
// Refit track only if it has an ultimate COT parent
// By default Mode=0, so far only StntupleEleFilter is initializing track block
// with Mode=1 to save time on the beam-constrained fit
//-----------------------------------------------------------------------------
  if (Mode == 0) {
    if (CotParent) {
      refit->beamcon(*CotParent);
      if (refit->ok()) {
	trk->fBcLam0  = refit->par()[CT_Standard::ict];
	trk->fBcC0    = refit->par()[CT_Standard::icu];
	trk->fBcZ0    = refit->par()[CT_Standard::iz0];
	trk->fBcPhi0  = refit->par()[CT_Standard::ip0];
	trk->fBcD0    = refit->par()[CT_Standard::id0];
      }
      else {
	//	logger->Report(-201,Form("StntupleInitTrack: CT_ReFit Failed %d ",
	//		 refit->status()));
	trk->fBcvind = refit->status() & 0xffff;
      }
    }
  }
  //  else  logger->Report(-103,"StntupleInitTrack: No CT_ReFit since not CotParent");
//-----------------------------------------------------------------------------
// Aaron: set the CdfTrack id number so that we can make the TStnTrackLinkBlock 
//        if we have to.  This variable is not streamed
// Pasha: I believe we don't need "fCdfTrackId" any longer, but having it 
// around may be convenient
//-----------------------------------------------------------------------------
  trk->fCdfTrackId = rick->id().value();
  trk->fCdfTrack   = (CdfTrack*) rick;

  trk->SetObspNumber(-1);
  trk->SetGenpNumber(-1);

  return 0;
}

//_____________________________________________________________________________
Int_t StntupleInitTrackBlock(TStnDataBlock* Block, 
			     AbsEvent*      Event, 
			     int            Mode) 
{
  // this routine is called from StntupleMakerModule (this assumption is 
  // essential) and fills track data block

  int                     ev_number, rn_number;

  static TFolder*         fol    = NULL;
  //  static TStnErrorLogger* logger = NULL;
  static StntupleCotReFit refit;
  static StntupleSiExpected siexpected;
//-----------------------------------------------------------------------------
//  initialize static variables
//-----------------------------------------------------------------------------
  if (! fol) {
    fol = (TFolder*) gROOT->GetRootFolder()->FindObject("StntupleFolder");
    //    logger     = (TStnErrorLogger) fol->FindObject("StntupleErrorLogger");
    fol->Add((TObject*) &refit);
    fol->Add((TObject*) &siexpected);
  }
					// check if block has already been 
					// initialized

  ev_number = AbsEnv::instance()->trigNumber();
  rn_number = AbsEnv::instance()->runNumber();

  if (Block->Initialized(ev_number,rn_number)) return 0;

  TStnTrackBlock* data = (TStnTrackBlock*) Block;
  data->Clear();

  CdfTrackView_h  theTracks;

  int rc = StntupleCreateTrackView(Event,data->CollName()->Data(),theTracks);
  if (rc < 0) return rc;

  const CdfTrackView::CollType& track_set = theTracks->contents();

  const CdfTrack* cdf_track, *ct1, *ct2;
  TStnTrack* trk;
  int ntracks = track_set.size();
//-----------------------------------------------------------------------------
// set up the collection of simulated particles (if there)
//-----------------------------------------------------------------------------
  EventRecord::ConstIterator obsp_iter(Event, "OBSP_StorableBank");
  ConstHandle<OBSP_StorableBank> obsp(obsp_iter);

  EventRecord::ConstIterator tmatch_iter(Event,"CdfTrackMatch");
  
  CdfTrackMatch* tmatch = 0;
  if (tmatch_iter.is_valid()) {
    ConstHandle<CdfTrackMatch> h_tmatch(tmatch_iter);
    tmatch = (CdfTrackMatch*) h_tmatch.operator->();
  }

  //-----------------------------------------------
  // setup stuff for siExpected
  //-----------------------------------------------

//-----------------------------------------------------------------------------
// order tracks in Pt
//-----------------------------------------------------------------------------
  int* inum = new int[ntracks];
  for (int i=0; i<ntracks; i++) inum[i] = i;

  for (int i1=0; i1<ntracks-1; i1++) {
    ct1 = (track_set)[inum[i1]].operator->();
    for (int i2=i1+1; i2<ntracks; i2++) {
      ct2 = (track_set)[inum[i2]].operator->();
      if (ct1->pt() < ct2->pt()) {
					// swap the tracks
	int itmp = inum[i1];
	inum[i1] = inum[i2];
	inum[i2] = itmp;
	ct1      = ct2;
      }
    }
  }
//-----------------------------------------------------------------------------
// tracks are ordered in Pt. If necessary, reinitialize refit object
//-----------------------------------------------------------------------------
  refit.BeginRun();
  siexpected.BeginRun();
  siexpected.BeginEvent(Event);
  for (int i=0; i<ntracks; i++) {
    int ind = inum[i];
    cdf_track = (track_set)[ind].operator->();
    trk       = data->NewTrack();
    StntupleInitTrack(trk,cdf_track,Mode);
  }
//-----------------------------------------------------------------------------
// calculate parameters which require track list to be defined
//-----------------------------------------------------------------------------
  for (int i=0; i<ntracks; i++) {
    trk       = data->Track(i);
    trk->SetIso4(TStntuple::TrkIso(trk,data->TrackList()));
//-----------------------------------------------------------------------------
// MC hit matching for each track, we may deal with the refitted tracks...
//-----------------------------------------------------------------------------
    if (tmatch) {
      const CdfTrack* prnt = trk->GetCdfTrack();
      while(prnt->parent().is_nonnull()) prnt = &*(prnt->parent());

      int iobsp = tmatch->match(prnt->id().value())-1;
      if ((iobsp >= 0) && obsp->is_valid()) {
	OBSP_StorableBank::particleIter p(*obsp,iobsp);
	trk->SetObspNumber(iobsp);
	trk->SetGenpNumber(obsp->HepgNumber(p));
      }
    }
  }

  data->f_EventNumber = ev_number;
  data->f_RunNumber   = rn_number;
					// delete tmp variables allocated on 
					// the heap
  delete [] inum;

  return 0;
}

//_____________________________________________________________________________
Int_t StntupleTrackBlockLinks(TStnDataBlock* Block, AbsEvent* Event, int Mode)
{
  // for each track define the number of the corresponding primary vertex
  // use Z-vertices - not sure what I'm doing!

  TStnVertexBlock* vblock;
  TStnEvent* ev;
  TStnTrackBlock* tblock = (TStnTrackBlock*) Block;
  int      best_v;
  double   dist, min_dist;

  ev     = Block->GetEvent();
  vblock = (TStnVertexBlock*) ev->GetDataBlock("ZVertexBlock");
    
  for(int i=0; i<tblock->NTracks(); i++) {
    TStnTrack* trk = tblock->Track(i);
    best_v   = -1;
    min_dist = 5.0;
    for(int iv=0; iv<vblock->NVertices(); iv++) {
      TStnVertex* v = vblock->Vertex(iv);
      dist = fabs(trk->Z0() - v->Z());
      if (dist < min_dist) {
        best_v   = iv;
        min_dist = dist;
      }          
    }
    trk->SetVertexNumber(best_v);
  }
  return 0;
}