#ifndef OSHKinFitter_h
#define OSHKinFitter_h

#include "OSUserHiggs/IOSHKinFitter.h"
#include "OSServices/IOSHistSvc.h"

#include "AthenaBaseComps/AthAlgTool.h"

#include "GaudiKernel/ToolHandle.h"

//#include "KinFitter/TFitParticlePtEtaPhi.h"
#include "KinFitter/TFitParticleRelE.h"
//#include "KinFitter/TFitConstraintM.h"
#include "KinFitter/TFitConstraintMBW.h"
#include "KinFitter/TFitConstraintMGaus.h"
//#include "KinFitter/TFitConstraintPt.h"
#include "KinFitter/TFitConstraintEp.h"
//#include "KinFitter/TFitConstraintEpGaus.h"
#include "KinFitter/TKinFitter.h"

#include <TLorentzVector.h>
#include <TMatrixD.h>

#include <string>
#include <vector>

namespace Analysis {

  /** 
      @author Carl Gwilliam <gwilliam@mail.cern.ch>
   */

  class OSHKinFitter: virtual public IOSHKinFitter, public AthAlgTool {
  public:
    
    OSHKinFitter(const std::string &type, const std::string &name, const IInterface *parent);
    virtual ~OSHKinFitter();
    
    StatusCode                         initialize();
    StatusCode                         finalize();    

    void                               fit();
    void                               reset();
    void                               setZll(HepLorentzVector& lep1, HepLorentzVector& lep2, int pdg, bool constrain);
    void                               setZjj(HepLorentzVector& jet1, HepLorentzVector& jet2, bool constrain);
    void                               setMETHZZllqq(std::vector<HepLorentzVector>& otherJets, float etMissX, float etMissY, bool constrain);

    TMatrixD*                          energyCovMatrix(double dE);
    TLorentzVector*                    hlvtotlv(HepLorentzVector& hlv);
    HepLorentzVector*                  tlvtohlv(TLorentzVector& tlv);
    void                               print();
    void                               deletePart(TAbsFitParticle*);


    HepLorentzVector*                  getZllLep1() {return tlvtohlv(*const_cast<TLorentzVector*>(m_zllLep1->getCurr4Vec()));}
    HepLorentzVector*                  getZllLep2() {return tlvtohlv(*const_cast<TLorentzVector*>(m_zllLep2->getCurr4Vec()));}

    HepLorentzVector*                  getZjjJet1() {return tlvtohlv(*const_cast<TLorentzVector*>(m_zjjJet1->getCurr4Vec()));}
    HepLorentzVector*                  getZjjJet2() {return tlvtohlv(*const_cast<TLorentzVector*>(m_zjjJet2->getCurr4Vec()));}

    std::vector<HepLorentzVector*>*    getOtherJets();

  private:

    TKinFitter                         m_fitter;
    int                                m_nIter;
    float                              m_maxDelta;
    float                              m_maxSum;
    float                              m_nFailed;

    TFitParticleRelE*                  m_zllLep1;
    TFitParticleRelE*                  m_zllLep2;
    TFitConstraintMBW*                 m_zllMassCons;

    TFitParticleRelE*                  m_zjjJet1;
    TFitParticleRelE*                  m_zjjJet2;
    TFitConstraintMBW*                 m_zjjMassCons;

    TFitConstraintEp*                  m_pxCons;
    TFitConstraintEp*                  m_pyCons;

    std::vector<TAbsFitParticle*>*     m_particles;
    std::vector<TAbsFitParticle*>*     m_otherJets;
  };
}

#endif