#include "TVector2.h"
#include <MetObjects/CdfMet.hh>

#include "CalorObjects/PhysicsTowerParams.hh"
#include "CalorObjects/PhysicsTowerData.hh"
#include "CalorObjects/TowerCorrectionColl.hh"
#include "Calor/PhysicsTowerDataMaker.hh"
#include "Calor/PhysicsTowerDataMaker.hh"

#include "MetObjects/CdfMet.hh"
#include "Met/NaiveEtParams.hh"
#include "Met/NaiveEtCalculator.hh"

#include <Stntuple/obj/TStnErrorLogger.hh>
#include <Stntuple/obj/TStnEvent.hh>
#include <Stntuple/obj/TStnVertexBlock.hh>
#include <Stntuple/obj/TStnTrackBlock.hh>
#include <Stntuple/obj/TStnLepton.hh>
#include <Stntuple/obj/TStnMetBlock.hh>
#include <Stntuple/obj/TStnDBManager.hh>

#include <Stntuple/alg/TStntuple.hh>

#include <Stntuple/mod/StntupleUtilities.hh>
#include <Stntuple/mod/InitStntupleDataBlocks.hh>


int StntupleCorrectMetForTracks(PhysicsTowerData_ch& TowerCol, 
				TStnTrackBlock*      TrackBlock,
				TVector2*            Met2      );

//_____________________________________________________________________________
Int_t StntupleInitMetBlock(TStnDataBlock* Block, 
                           AbsEvent*      Event,
                           int            Mode) 
{
					// fill MET block with the `event' data
  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;

  TStnMetBlock* data = (TStnMetBlock*) Block;
  data->Clear();
				// Here, grab the collection by name:

  StntupleGetProcessName(data->CollName()->Data(),process,description); 

  CdfMet_ch met;
  CdfMet::Error metStatus;

  if (strcmp(description,"none") != 0) {
     metStatus = CdfMet::find(met,description);
  }
  else {
    metStatus = CdfMet::find(met);
  }
  
  if (metStatus == CdfMet::ERROR) {
    TStnEvent* ev = data->GetEvent();
    TStnErrorLogger* logger = ev->GetErrorLogger();
    logger->Report(-101,Form("InitMetBlock: cant find CdfMet for %s:%s",
			     process,
			     description));
    return -101;
  }
  
  data->fMet   [0]= met->met();
  data->fMetphi[0]= TVector2::Phi_0_2pi(atan2(-met->eySum(),-met->exSum()));
  data->fMetjt    = TStnDataBlock::kUndefined;
  data->fMetjtphi = TStnDataBlock::kUndefined;
  data->fMetj1    = TStnDataBlock::kUndefined;
  data->fMetj1phi = TStnDataBlock::kUndefined;
  data->fMetc     = TStnDataBlock::kUndefined;
  data->fMetcphi  = TStnDataBlock::kUndefined;
  data->fHtc      = TStnDataBlock::kUndefined;
  data->fMetcf    = TStnDataBlock::kUndefined;
  data->fMetcfphi = TStnDataBlock::kUndefined;
  data->fHtcf     = TStnDataBlock::kUndefined;
  data->fEtatrks  = TStnDataBlock::kUndefined;
  data->fPhitrks  = TStnDataBlock::kUndefined;
  data->fPttrks   = TStnDataBlock::kUndefined;
  data->fSeedpt   = TStnDataBlock::kUndefined;
  data->fSumet    = met->etSum();
  data->fSumetjet = TStnDataBlock::kUndefined;
  data->fMetsig   = TStnDataBlock::kUndefined;
//-----------------------------------------------------------------------------
//  set met[1]-met[4] to met[0] (this is the default)
//-----------------------------------------------------------------------------
  for (int i=1; i<5; i++) {
    data->fMet   [i] = data->fMet   [0];
    data->fMetphi[i] = data->fMetphi[0];
  }
//-----------------------------------------------------------------------------
//  on return mark block as initialized for given event/run
//-----------------------------------------------------------------------------
  data->f_RunNumber   = rn_number;
  data->f_EventNumber = ev_number;

  return 0;
}


//_____________________________________________________________________________
Int_t StntupleMetBlockLinks(TStnDataBlock* Block, AbsEvent* Event, int Mode) 
{

  static PhysicsTowerDataMaker ptd_maker;
  static NaiveEtCalculator     met_calculator;

  bool                         status;
  CdfMet_h                     metz;
  PhysicsTowerData_ch          towers_ch;

                                        // We would like to recalculate Met 
                                        // at z vertex from Beate's ZVertexBlock 

  static PhysicsTowerDataMaker ptd_maker_dt;
  static NaiveEtCalculator     met_calculator_dt;
  PhysicsTowerData_ch          towers_ch_dt;

  Int_t                        rc, ev_number, rn_number, nlep;
  Float_t                      sumpt;
  Double_t                     vz(0);
                                       // parameters for 3-D vertex calculation

  static PhysicsTowerDataMaker ptd_maker_xy;
  static NaiveEtCalculator     met_calculator_xy;
  PhysicsTowerData_ch          towers_ch_xy;

  TStnMetBlock* met_block = (TStnMetBlock*) Block;

  TStnVertexBlock* vtx_block;
                                        // pointer to z vertex block
                                        // to recalculate met at highest sum pt z vertex
  TStnVertexBlock* zvtx_block;           

  TStnTrackBlock*  trk_block;
  TObjArray*       hptl;
					// assume that the block exists, it 
					// is supposed to be initialized

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

  TStnEvent* ev = met_block->GetEvent();

  if (! Block->Initialized(ev_number,rn_number)) {
    TStnErrorLogger* logger = ev->GetErrorLogger();
    logger->Report(-102,"StntupleMetBlockLinks: MetBlock not initialized ");
    return -1;
  }
					// do not do initialize links 2nd time

  if (Block->LinksInitialized()) return 0;

  rc = 0;
  trk_block = (TStnTrackBlock*) ev->GetDataBlock("TrackBlock");
//-----------------------------------------------------------------------------
// decide on what we call the event vertex:
// - if there are high-Pt lepton(s) in the event, use Z0 of teh 1st one, 
//   otherwise take the first vertex from TStnVertexBlock
// "StntupleHptl" added in InitStntupleModule
//-----------------------------------------------------------------------------
  hptl = ev->GetListOfHptl();
  nlep = hptl->GetEntriesFast();
  if (nlep > 0) {
    TStnLepton* lep = (TStnLepton*) hptl->UncheckedAt(0);
    vz = TStntuple::LeptonZ0(lep,trk_block);
  }
  else {
    vtx_block = (TStnVertexBlock*) ev->GetDataBlock("VertexBlock");
    if (vtx_block->NVertices() > 0) {
//-----------------------------------------------------------------------------
//  take Z coordinate of the 1st vertex in "VertexBlock"
//-----------------------------------------------------------------------------
      TStnVertex* v = vtx_block->Vertex(0);
      vz = v->Z();
    }
  }
//-----------------------------------------------------------------------------
// met[1]: MET in the 1st vertex from VertexBlock p2=2: MetCalculator
//-----------------------------------------------------------------------------
  PhysicsTowerParams pdt_parms(0,2,0.1,vz);
  towers_ch = ptd_maker.create(Event,pdt_parms,"StntupleFilterModule");
  if (vz != 0) {
    int vtx_strategy = 1;
    NaiveEtParams met_parms(0.1,vtx_strategy);
    metz                  = new CdfMet(vz,vtx_strategy);
    status                = met_calculator.calculateMet(metz,towers_ch);
    met_block->fMet   [1] = metz->met();
    met_block->fMetphi[1] = TVector2::Phi_0_2pi(atan2(-metz->eySum(),
						      -metz->exSum()));
    met_block->fHtc = metz->etSum();
  }
//-----------------------------------------------------------------------------
// met[2]: corrected for the tracks
//-----------------------------------------------------------------------------
  TVector2 met2;
  StntupleCorrectMetForTracks(towers_ch,trk_block,&met2);
  met_block->fMet   [2] = met2.Mod();
  met_block->fMetphi[2] = TVector2::Phi_0_2pi(met2.Phi());
//-----------------------------------------------------------------------------
//  met[3]: in the highest-Pt Zvertex
//-----------------------------------------------------------------------------
  sumpt = 0;
  zvtx_block  = (TStnVertexBlock*) ev->GetDataBlock("ZVertexBlock");
  for (int i = 0; i < zvtx_block->NVertices(); i++) {
    TStnVertex* v = zvtx_block->Vertex(i);
    if ( v->SumPt() > sumpt) {
      sumpt = v->SumPt();
      vz    = v->Z(); 
    }
  }

  if (vz != 0) {
    PhysicsTowerParams pdt_parms_dt(0,2,0.1,vz);
    towers_ch_dt = ptd_maker_dt.create(Event,pdt_parms_dt,"StntupleFilterModule");
    int vtx_strategy = 1;
    NaiveEtParams met_parms(0.1,vtx_strategy);
    metz                  = new CdfMet(vz,vtx_strategy);
    status                = met_calculator_dt.calculateMet(metz,towers_ch_dt);
    met_block->fMet   [3] = metz->met();
    met_block->fMetphi[3] = TVector2::Phi_0_2pi(atan2(-metz->eySum(),
						      -metz->exSum()));
    met_block->fSumetjet = metz->etSum();
  }
//-----------------------------------------------------------------------------
//  met[4]: take Z coordinate of the highest sum pt vertex (vz, found above)
//          and the x and y from the SVX beamline
//-----------------------------------------------------------------------------
  TStnDBManager* dbm = TStnDBManager::Instance();
  TStnBeamPos* svx_beam_pos = (TStnBeamPos*) dbm->GetTable("SvxBeamPos");
  float threeVertex[3];
  float slope[2];

  if (svx_beam_pos && svx_beam_pos->Status()==0 ) {
    slope      [0] = svx_beam_pos->DxDz();
    slope      [1] = svx_beam_pos->DyDz();
    threeVertex[0] = svx_beam_pos->X0() + vz*slope[0];
    threeVertex[1] = svx_beam_pos->Y0() + vz*slope[1];
  } else {
    for(int i=0; i<3; i++) threeVertex[i] = 0.0;    
    for(int i=0; i<2; i++) slope[i] = 0.0;
  }

  if (vz != 0) {
    threeVertex[2] = vz;
    PhysicsTowerParams pdt_parms_xy(0,2,0.1,threeVertex,slope);
    towers_ch_xy   = ptd_maker_xy.create(Event,pdt_parms_xy,"StntupleFilterModule");
    metz                  = new CdfMet(threeVertex, slope, 1);
    status                = met_calculator_xy.calculateMet(metz,towers_ch_xy);
    met_block->fMet   [4] = metz->met();
    met_block->fMetphi[4] = TVector2::Phi_0_2pi(atan2(-metz->eySum(),
						      -metz->exSum()));
  }
//-----------------------------------------------------------------------------
// mark links as initialized
//-----------------------------------------------------------------------------
  met_block->fLinksInitialized = 1;

  return rc;
}

//_____________________________________________________________________________
int StntupleCorrectMetForTracks(PhysicsTowerData_ch& TowerCol, 
				TStnTrackBlock*      TrackBlock,
				TVector2*            Met2      ) 
{
  // assume that all the tracks are well measured....

  Double_t    phi, pt_cal, pt_trk, chi2_cot;
  Int_t       wedge, nt, track_is_good;

  TVector2    mom24_cal[24];
  TVector2    mom24_trk[24];

  TVector2    mom;

  PhysicsTowerData::const_iterator end = TowerCol->end();
//-----------------------------------------------------------------------------
//  loop over the towers and calculate momentum vector for each wedge, 
//  assuming 24-fold symmetry
//-----------------------------------------------------------------------------
  for (int i=0; i<24; i++) {
    mom24_cal[i].Set(0.,0.); 
    mom24_trk[i].Set(0.,0.); 
  }

  for (PhysicsTowerData::const_iterator iter=TowerCol->begin(); iter != end; ++iter){
    PhysicsTower* twr = iter.tower();
    mom.Set(twr->totEt()*cos(twr->totPhi()),twr->totEt()*sin(twr->totPhi()));
    phi               = TVector2::Phi_0_2pi(mom.Phi());
    wedge             = int( phi/(TMath::Pi()/12.) );
    mom24_cal[wedge] += mom;
  }
//-----------------------------------------------------------------------------
//  done with the towers, loop over the tracks and do the same
//-----------------------------------------------------------------------------
  nt = TrackBlock->NTracks();
  for (int i=0; i<nt; i++) {
    TStnTrack* trk = TrackBlock->Track(i);
    mom.Set(trk->Momentum()->Px(),trk->Momentum()->Py());
//-----------------------------------------------------------------------------
//  make sure that the track is reliable, for high-Pt tracks require 4 axial
//  segments to be available (the selection can and should be improved)
//-----------------------------------------------------------------------------
    track_is_good = 0;
    if (trk->NStSeg() < 2)                                  goto NEXT_TRACK;
    chi2_cot = trk->Chi2Cot()/(trk->NCotHitsTot()-4.9999);
    if (chi2_cot > 3.)                                      goto NEXT_TRACK;

    if (trk->Momentum()->Pt() < 40) {
      if (trk->NAxSeg() < 3)                                goto NEXT_TRACK;
    }
    else {
      if (trk->NAxSeg() < 4)                                goto NEXT_TRACK;
    }
//-----------------------------------------------------------------------------
// track passed the quality cuts, count it
//-----------------------------------------------------------------------------
    phi = TVector2::Phi_0_2pi(mom.Phi());
    wedge             = int (phi/(TMath::Pi()/12.));
    mom24_trk[wedge] += mom;

    NEXT_TRACK:;
  }
//-----------------------------------------------------------------------------
//  consistency check: met1 should be the same as met, corrected for the Z,
//  met2 can be better (if all the tracks are OK)
//-----------------------------------------------------------------------------
  TVector2 met1(0,0);
  Met2->Set(0.,0);

  for (int iw=0; iw<24; iw++) {
    pt_cal = mom24_cal[iw].Mod();
    pt_trk = mom24_trk[iw].Mod();
    met1 += mom24_cal[iw];
    if (pt_trk > pt_cal) {
//-----------------------------------------------------------------------------
//  assume that track went into a crack... and do the best
//-----------------------------------------------------------------------------
      *Met2 += mom24_trk[iw];
    }
    else {
      *Met2 += mom24_cal[iw];
    }
  }
  return 0;
}