///////////////////////////////////////////////////////////////////////////////
// P.Murat
///////////////////////////////////////////////////////////////////////////////
#include "TROOT.h"
#include "TFolder.h"

#include <BaBar/Experiment.hh>
#include <Edm/EventRecord.hh>
#include <Edm/ConstHandle.hh>

#include <TrackingObjects/Storable/CdfTrackView.hh>

#include "Calor/CalorPredicates.hh"
#include "Calor/PhysicsTowerDataMaker.hh"
#include "Calor/calor_alg.hh"
#include "Calor/calorMakeView.hh"
#include "CalorGeometry/CalParameters.hh"
#include "CalorObjects/PhysicsTowerView.hh"

#include <TauObjects/Pi0Candidate.hh>
#include <TauObjects/Pi0CandidateView.hh>
#include <TauObjects/TauCluster.hh>
#include <TauObjects/CdfTau.hh>
#include <TauObjects/CdfTauCollection.hh>
#include "Tau/algorithms/TauClusterRenewCache.hh"

#include <Stntuple/alg/TStntuple.hh>

#include <Stntuple/mod/InitStntupleDataBlocks.hh>
#include <Stntuple/mod/StntupleUtilities.hh>
#include <Stntuple/obj/TStnEvent.hh>
#include <Stntuple/obj/TStnErrorLogger.hh>
#include <Stntuple/obj/TStnTauBlock.hh>
#include <Stntuple/obj/TStnTrackBlock.hh>
#include <Stntuple/obj/TStnPi0Block.hh>
#include <Stntuple/obj/TObspBlock.hh>

namespace {
  PhysicsTowerData_ch    fgTauPhysicsTowers;
}

//_____________________________________________________________________________
Int_t StntupleInitTau(TStnTau* Tau, const CdfTau* Fedor, Int_t Itau) {
  // init TStnTau from CdfTau (Run II only)

  TauCluster* fedors_cluster = (TauCluster*) Fedor->caloCluster();

  Tau->fNumber      = Itau;
  Tau->fCdfTau      = Fedor;
  Tau->fDecayMode   = TStnTau::kHadronicMode;
//-----------------------------------------------------------------------------
// for the time being the best momentum is the one measured by the calorimeter
//-----------------------------------------------------------------------------
  HepLorentzVector pi0_mom;
  HepLorentzVector trk_mom;

  Fedor->tracksFourMomentum(&trk_mom);
  Fedor->pi0sFourMomentum  (&pi0_mom);

  Tau->fTrackMomentum.SetXYZT(trk_mom.x(),
			      trk_mom.y(),
			      trk_mom.z(),
			      trk_mom.e());

  Tau->fVisMomentum.SetXYZT(trk_mom.x()+pi0_mom.x(),
			    trk_mom.y()+pi0_mom.y(),
			    trk_mom.z()+pi0_mom.z(),
			    trk_mom.e()+pi0_mom.e());

  Tau->fClusterMomentum.SetXYZT(Fedor->clusterPx(), 
				Fedor->clusterPy(),
				Fedor->clusterPz(),
				Fedor->clusterE());


  Tau->fMomentum.SetXYZT(Fedor->clusterPx(), 
			 Fedor->clusterPy(),
			 Fedor->clusterPz(),
			 Fedor->clusterE());

  Tau->fGenpMomentum.SetXYZT(0.,0.,0.,1.);

  Tau->fCharge      = Fedor->tracksChargeInSeed10Degree();

  Tau->fDteta       = Fedor->clusterDetectorEta();
  Tau->fZv          = Fedor->seedTrackVertexZ();
					// not sure what it should be

  Tau->fStatusCode  = TStnDataBlock::kUndefined;
//-----------------------------------------------------------------------------
// energy-weighted cluster momenta
//-----------------------------------------------------------------------------
  Tau->fEtaEta      = Fedor->clusterEtaEtaMoment();
  Tau->fPhiPhi      = Fedor->clusterPhiPhiMoment();
  //Tau->fDelr        = sqrt(Tau->fEtaEta*Tau->fEtaEta+Tau->fPhiPhi*Tau->fPhiPhi);

  Tau->fEtEm        = Fedor->clusterEmEt();
  Tau->fEtHad       = Fedor->clusterHadronEt();
  Tau->fEmfr        = Fedor->emFraction(); //Fedor->clusterEmEt()/Fedor->clusterEt();
  Tau->fNMuStubs    = Fedor->muonStubsNumber();
  Tau->fNMuHits     = Fedor->muonHitsNumber();


  Tau->fNTracks10    = Fedor->chargedTracksNumberSeed10Degree();
  Tau->fNTracks30    = Tau->fNTracks10 + Fedor->chargedTracksNumberSeed10To30Degree(1.);
  Tau->fNWrongTracks = Fedor->wrongVertexTracksNumberSeed10Degree();

  Tau->fNExtrapTracks   = SHRT_MAX;
  Tau->fNExtrapTracks10 = SHRT_MAX;
  Tau->fNExtrapTracks30 = SHRT_MAX;

  Tau->fClslNumber   = -1;
  Tau->fJetsNumber   = -1;
  Tau->fTauoNumber   = -1;

  Tau->fNL2Towers    = 20;
  Tau->fNTowers      = Fedor->nTowers();
//-----------------------------------------------------------------------------
// order towers in Et
//-----------------------------------------------------------------------------
  PhysicsTower*  tow[100];
  TowerKey       tau_key[100];

  vector<const PhysicsTower*>& tower_list = fedors_cluster->towers().contents();
  int ntowers = tower_list.size();
  for (int i=0; i<ntowers; i++) {
    tow[i] = (PhysicsTower*) tower_list[i];
  }
	
  PhysicsTower* buf;
  for (int i=0; i<ntowers-1; i++) {
    double eti = tow[i]->totEt();
    for (int j=i+1;j<ntowers; j++) {
      double etj = tow[j]->totEt();
      if (eti < etj) {
	buf    = tow[i];
	tow[i] = tow[j];
	tow[j] = buf;
	eti    = etj;
      }
    }
  }
//-----------------------------------------------------------------------------
// store keys for ordered towers
//-----------------------------------------------------------------------------
  for (int i=0; i<ntowers; i++) {
    tau_key[i] = tow[i]->key();
  }
					// should be at least one tower
  Tau->fTowerEt[0]   = 0;
  Tau->fTowerEt[1]   = 0;

  if (ntowers > 0) Tau->fTowerEt[0]   = tow[0]->totEt();
  if (ntowers > 1) Tau->fTowerEt[1]   = tow[1]->totEt();

  Tau->fNPi010       = Fedor->pi0sNumber();
  Tau->fNPi01030     = Fedor->isolationPi0s(0.5);

  Tau->fSumPt30      = Fedor->trackIsolation ();
  Tau->fSumPtPi030   = Fedor->isoPi0sScalarPt ();
  
//-----------------------------------------------------------------------------
// seed track parameters
//-----------------------------------------------------------------------------
  Tau->fSeedTrackZ   = Fedor->seedTrackVertexZ();
  Tau->fSeedTrackPt  = Fedor->seedTrackPt();
  Tau->fSeedDeltaPhi = Fedor->phiAngleSeedTrackToCluster();

  Tau->fSumTrackP    = Fedor->tracksScalarPSum();
  Tau->fSumPt10      = Fedor->tracksScalarPtSum();;
  Tau->fNTrkTau      = 20;

  Tau->fECorrectionFactor = Fedor->eCorrectionFactor();
  Tau->fTauCone           = Fedor->tauCone();
  Tau->fTauConePi0        = Fedor->tauConePi0();

/*
  PhysicsTowerView R4Towers;

  Double_t eta1 = Fedor->caloCluster()->eta();
  Double_t eta  = Fedor->caloCluster()->etaDetector();
  Double_t phi  = TVector2::Phi_0_2pi(Fedor->caloCluster()->phi());

  calorMakeView(fgTauPhysicsTowers,R4Towers,NeighborsInCone(eta,phi,0.4));

  Towers::const_iterator firstTower = R4Towers.begin();
  Towers::const_iterator lastTower  = R4Towers.end();

  double em04  = 0;
  double had04 = 0;
  for (Towers::const_iterator iter = firstTower; iter != lastTower; ++iter) {
    if (*iter) {
//-----------------------------------------------------------------------------
// make sure we are not summing over the tau towers
//-----------------------------------------------------------------------------
      int found = 0;
      TowerKey key = (*iter)->key();
      for (int i=0; i<ntowers; i++) {
	if (key == tow[i]->key()) {
	  found = 1;
	  break;
	}
      }
      if (!found) {
	had04 += (*iter)->hadEt();
	em04  += (*iter)->emEt();
      }
    }
  }

  Tau->fEmIso04      = em04;
  Tau->fHadIso04     = had04;
  Tau->fCalIso04     = em04+had04;
*/

//-----------------------------------------------------------------------------
// note: in CdfTau calor isolation is in DR<0.4 cone around the extrapolated seed
//-----------------------------------------------------------------------------
  Tau->fEmIso04      = Fedor->caloIsolationEmEt();
  Tau->fHadIso04     = Fedor->caloIsolationHadronEt();
  //Tau->fCalIso04     = em04+had04;
  Tau->fTrackIso     = TStnDataBlock::kUndefined;

  return 0;
}


//_____________________________________________________________________________
Int_t StntupleInitTauBlock(TStnDataBlock* block, AbsEvent* event, int mode) 
{
  // initialize tau data block with the `event' data
  // return -1, if bank doesn't exist, 0, if everything is OK

				// in hope that number of tau candidates is 
				// always less than 100
  int   inum[100];
  int   ev_number, rn_number;
  char  process[100], description[100];
//-----------------------------------------------------------------------------
//  save time: don't do initialization twice for the same event
//-----------------------------------------------------------------------------
  ev_number = AbsEnv::instance()->trigNumber();
  rn_number = AbsEnv::instance()->runNumber();

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


  TStnTauBlock* data = (TStnTauBlock*) block;
  data->Clear();

  StntupleGetProcessName(data->CollName()->Data(),process,description);
  ConstHandle<CdfTauCollection> handle (StntupleGetIterator(event,
							    "CdfTauCollection",
							    process,
							    description));
  if (handle.is_null()) return -1;

  const CdfTauCollection::CollType& tau_list = handle->contents(); 
  int   ntaus = tau_list.size();

  if (ntaus == 0) return 0;

  const CdfTau*          fedor;
  TStnTau*               tau;
//-----------------------------------------------------------------------------
// convert CalData into a list of "physics towers" to calculate isolation
//-----------------------------------------------------------------------------
  PhysicsTowerDataMaker  maker;
  PhysicsTowerParams     params(0,0,0.,0.);

  fgTauPhysicsTowers  = maker.create(event,params,"TauTowers");
//-----------------------------------------------------------------------------
// order taus in Et
//-----------------------------------------------------------------------------
  const CdfTau *tau1, *tau2;

  for (int i=0; i<ntaus; i++) inum[i] = i;

  for (int i1=0; i1<ntaus-1; i1++) {
    tau1 = &tau_list[i1];
    for (int i2=i1+1; i2<ntaus; i2++) {
      tau2 = &tau_list[i2];
      if (tau1->clusterEt() < tau2->clusterEt()) {

					// swap the taus
	int itmp = inum[i1];
	inum[i1] = inum[i2];
	inum[i2] = itmp;
	tau1     = tau2;
      }
    }
  }

  for (int i=0; i<ntaus; i++) {
    fedor = &tau_list[inum[i]];
//-----------------------------------------------------------------------------
// make sure tau cluster information has been updated 
// if tau reco code has been ran in this job, it would not be necessary, 
// but should work anyway
//-----------------------------------------------------------------------------
    TauClusterFriend::renewCache (fedor->caloCluster(), 
 				  fedor->referenceVertexZ());
    tau  = data->NewTau();
    StntupleInitTau(tau,fedor,i);
  }
//-----------------------------------------------------------------------------
//  on return mark block as initialized for given event/run
//-----------------------------------------------------------------------------
  data->f_RunNumber   = rn_number;
  data->f_EventNumber = ev_number;

  return 0;
}



//_____________________________________________________________________________
Int_t StntupleTauBlockLinks(TStnDataBlock* block, AbsEvent* event, int mode) 
{
  static TFolder*         fol       = NULL;
  static TStnErrorLogger* logger    = NULL;

  // set links between the tau's and the objects in the other blocks
  // return 0 if everything is OK

  int                ev_number, rn_number, npi0, ia;
//-----------------------------------------------------------------------------
// make sure the block is initialized, do not initialize links twice
//-----------------------------------------------------------------------------
  ev_number = AbsEnv::instance()->trigNumber();
  rn_number = AbsEnv::instance()->runNumber();

  if (! block->Initialized(ev_number,rn_number)) return -1;
  if (block->LinksInitialized()                ) return  0;

  if (! fol) {
//-----------------------------------------------------------------------------
//  initialize local static variables
//-----------------------------------------------------------------------------
    fol     = (TFolder*) gROOT->GetRootFolder()->FindObject("StntupleFolder");
    logger  = (TStnErrorLogger*) fol->FindObject("StntupleErrorLogger");
  }

  TStnTauBlock* tau_block = (TStnTauBlock*) block;

  const CdfTau*                             fedor;
  const vector<CdfTrack_clnk>*              track_list;
  const vector<const PhysicsTower*>*        tower_list;
  const Pi0CandidateView::collection_type * pi0_list;

  TStnTau*        tau;
  TStnEvent*      ev;
  TStnTrackBlock* trk_block;
  TStnPi0Block*   pi0_block;
  TStnTrack*      seed_track;
  Int_t           ntrk, ntowers, found;
//-----------------------------------------------------------------------------
// need pointer to the default track block
//-----------------------------------------------------------------------------
  ev        = tau_block->GetEvent();
  trk_block = (TStnTrackBlock*) ev->GetDataBlock("TrackBlock");
  pi0_block = (TStnPi0Block*  ) ev->GetDataBlock("Pi0Block"  );

  ntrk      = trk_block->NTracks();

  TObjArray* hptl = ev->GetListOfHptl();

  TVector3 v3;
  TVector3 seed_v3;
  int      iseed, itrk, nt_tau;
  double   angle, pt, trk_iso, sumpt;

  Int_t ntaus = tau_block->NTaus();
//-----------------------------------------------------------------------------
// make sure link lists are always initialized - is this correct
//-----------------------------------------------------------------------------
  tau_block->TrackLinkList()->Init(ntaus);
  tau_block->TowerLinkList()->Init(ntaus);
  tau_block->Pi0LinkList  ()->Init(ntaus);

  for (int itau=0; itau<ntaus; itau++) {
    tau        = tau_block->Tau(itau);
    fedor      = tau->GetCdfTau();
//-----------------------------------------------------------------------------
// fill track link block
//-----------------------------------------------------------------------------
// tau's 10deg prongs first:
    track_list = &fedor->contributedTauTracks()->contents();
    nt_tau = track_list->size();
    for (itrk=0; itrk<nt_tau; itrk++) {
      const CdfTrack* trk =  (*track_list)[itrk].operator->();
      int it = trk_block->TrackNumber(trk);
      tau_block->TrackLinkList()->Add(itau,it);
    }
// now ISO tracks:
    track_list = &fedor->isolationTauTracks()->contents();
    nt_tau = track_list->size();
    for (itrk=0; itrk<nt_tau; itrk++) {
      const CdfTrack* trk =  (*track_list)[itrk].operator->();
      int it = trk_block->TrackNumber(trk);
      tau_block->TrackLinkList()->Add(itau,it);
    }

    iseed      = tau_block->TrackNumber(itau,0);
    if (iseed < 0) {
//-----------------------------------------------------------------------------
// let the code crash at this point - this should not be happening, 
// so I want to fix the problem instead of complaining about it
//-----------------------------------------------------------------------------
      seed_track = 0;
    }
    else {
      seed_track = trk_block->Track(iseed);
    }
    if (seed_track == 0) {
      logger->Report(-104,
		     Form("StntupleTauBlockLinks: tau no. %i w/o seed track",
			  itau));
      goto NEXT_TAU;
    }
    seed_v3 = seed_track->Momentum()->Vect();
//-----------------------------------------------------------------------------
// recalculate tracking isolation - use all the tracks originating from the 
// same vertex (probably need to check that the tracks originate from the same 
// vertex with tau)
//-----------------------------------------------------------------------------
    sumpt = 0;
    for (int i=0; i<ntrk; i++) {
      TStnTrack* trk = trk_block->Track(i);
      if (trk->Momentum()->Pt() > 0.4) {
	found = 0;
	for (int itt=0; itt<tau->NTracks10(); itt++) {
	  itrk = tau_block->TrackNumber(itau,itt);
	  if (itrk == i) {
	    found = 1;
	    break;
	  }
	}
	if (! found) {
	  v3    = trk->Momentum()->Vect();
	  angle = seed_v3.Angle(v3)*180./TMath::Pi();
	  if ( (angle > 10) && (angle < 30) ) {
	    pt = trk->Momentum()->Pt();
	    sumpt += pt;
	  }
	}
      }
    }

    trk_iso = sumpt/tau->TrackMomentum()->Pt();
    tau->SetTrackIso(trk_iso);
//-----------------------------------------------------------------------------
// fill tower link block
//-----------------------------------------------------------------------------
    tower_list = &fedor->caloCluster()->towers().contents();
    ntowers = tower_list->size();
    for (int itower=0; itower<ntowers; itower++) {
      const PhysicsTower* tower =  (*tower_list)[itower];

      int tower_key = (((tower->iPhi() & 0x3f  )      ) |
		       ((tower->iEta() & 0x3f  ) <<  8) |
		       ((tower->key().isValid()) << 16)    );

      tau_block->TowerLinkList()->Add(itau,tower_key);
    }
//-----------------------------------------------------------------------------
// fill pi0 link block - protect against Pi0 block not being filled
//-----------------------------------------------------------------------------
    if (pi0_block) {
// tau's pi0's first:
      if (fedor->tauPi0s() != 0) {
	pi0_list = &fedor->tauPi0s()->contents();
	npi0 = pi0_list->size();
	for (int i=0; i<npi0; i++) {
	  const Pi0Candidate* anton =  &(*(*pi0_list)[i]);
	  ia = pi0_block->Pi0Number(anton);
	  if (ia >= 0)  tau_block->Pi0LinkList()->Add(itau,ia);
	  else {                 // this should not happen
	    logger->Report(-104, 
                  Form("StntupleTauBlockLinks: tau no. %i bad Pi0", itau));
	  }
	}
      }
// tau's ISO pi0's next:
      if (fedor->isoPi0s() != 0) {
	pi0_list = &fedor->isoPi0s()->contents();
	npi0 = pi0_list->size();
	for (int i=0; i<npi0; i++) {
	  const Pi0Candidate* anton =  &(*(*pi0_list)[i]);
	  ia = pi0_block->Pi0Number(anton);
	  if (ia >= 0) tau_block->Pi0LinkList()->Add(itau,ia);
	  else {                 // this should not happen
	    logger->Report(-104, 
                  Form("StntupleTauBlockLinks: tau no. %i bad ISO Pi0", itau));
	  }
	}
      }
    }
//-----------------------------------------------------------------------------
//  update list of good HPT leptons to calculate MET (if any...) - 
//  we do have isolation at this point
//-----------------------------------------------------------------------------
    if (TStntuple::GoodHptTau(tau)) {
      hptl->Add(tau);
    }
  NEXT_TAU:;
  }
//-----------------------------------------------------------------------------
// for MC data try to find a MC parent of a reconstructed tau. Do it for OBSP
// only - the procedure should be unambigous
//-----------------------------------------------------------------------------
  if (AbsEnv::instance()->monteFlag()) {
    TObspBlock* obsp;
    obsp = (TObspBlock*) tau_block->GetEvent()->GetDataBlock("ObspBlock");
    if (obsp) {
      TObjArray mc_list;
					// find list of OBSP taus, both tau+ 
					// and tau- are reported
      double dr, dr_min;
      int    nmc, best;

      nmc = obsp->GetParticles(15,&mc_list);

      if (nmc) {
	for (int i=0; i<ntaus; i++) {
	  TStnTau* reco_tau = tau_block->Tau(i);
	  dr_min = 0.3;
	  best   = -1;
	  for (int j=0; j<nmc; j++) {
	    TObspParticle* mc_tau = obsp->Particle(j);
	    dr = reco_tau->Momentum()->DrEtaPhi(*mc_tau->Momentum());
	    if (dr < dr_min) {
	      best   = mc_tau->Number();
	      dr_min = dr;
	    }
	  }
	  reco_tau->SetObspNumber(best);
	}
	mc_list.Clear();
      }
    }
  }
//-----------------------------------------------------------------------------
// mark links as initialized
//-----------------------------------------------------------------------------
  block->SetLinksInitialized();

  return 0;
}