#include "OSHKinFitter.h"
#include "CLHEP/Units/SystemOfUnits.h"
//#include "GaudiKernel/SystemOfUnits.h"

#include "boost/foreach.hpp"

#include <sstream>

static const double Zmass = 91.1876 * 1000;
static const double Zwidth = 2.4952 * 1000;

namespace Analysis {

  OSHKinFitter::OSHKinFitter(const std::string& type, const std::string &name, const IInterface *parent) : 
    AthAlgTool(type, name, parent),
    m_fitter(),
    m_nFailed(0),
    m_zllLep1(0),
    m_zllLep2(0),
    m_zllMassCons(0),
    m_zjjJet1(0),
    m_zjjJet2(0),
    m_zjjMassCons(0),
    m_pxCons(0),
    m_pyCons(0),
    m_particles(0),
    m_otherJets(0)
  {
    declareInterface<IOSHKinFitter>(this);

    declareProperty("Iterations",                m_nIter    = 3000                     );
    declareProperty("MaxDelta",                  m_maxDelta = 5e-3                     );
    declareProperty("MaxSumConstraints",         m_maxSum   = 1e-4                     );
  }


  OSHKinFitter::~OSHKinFitter() {
    // Destructor

  }  
  
  StatusCode OSHKinFitter::initialize() {
    // Initialize
    msg(MSG::INFO) << "OSHKinFitter::initialize()" << endreq;

    // Set up fitter
    m_fitter.setMaxNbIter(m_nIter);      // number of maximal iterations
    m_fitter.setMaxDeltaS(m_maxDelta);   // max Delta Chi2 between n and n-1 iteration
    m_fitter.setMaxF( 1e-4 );            // max sum of constraints
    // verbosity level // Verbosty of the fitter 0: quiet, 1: print result, 2: print iterations, 3: print also matrices
    m_fitter.setVerbosity(msgLvl(MSG::VERBOSE) ? 3 : (msgLvl(MSG::DEBUG) ? 2 : (msgLvl(MSG::INFO) ? 1 : 0 )));  

    // Init list of all particles
    m_particles = new std::vector<TAbsFitParticle*>();
    m_particles->reserve(8);

    // Init list of jets
    m_otherJets = new std::vector<TAbsFitParticle*>();

    return StatusCode::SUCCESS;
  }

  StatusCode OSHKinFitter::finalize() {
    // Finalize    
    msg(MSG::INFO) << "OSHKinFitter::Finalize()" << endreq;
    msg(MSG::INFO) << "Failed Fits = " << m_nFailed << endreq;
    return StatusCode::SUCCESS;
  }

  void OSHKinFitter::reset() {
    // Need to free memory - also for vec and cov as not owned by fit clases

    // Is this enough or do we need to renew each time?
    m_fitter.reset();
    m_fitter.setMaxNbIter(m_nIter);      // number of maximal iterations
    m_fitter.setMaxDeltaS(m_maxDelta);   // max Delta Chi2 between n and n-1 iteration
    m_fitter.setMaxF( 1e-4 );            // max sum of constraints
    // verbosity level // Verbosty of the fitter 0: quiet, 1: print result, 2: print iterations, 3: print also matrices
    m_fitter.setVerbosity(msgLvl(MSG::VERBOSE) ? 3 : (msgLvl(MSG::DEBUG) ? 2 : (msgLvl(MSG::INFO) ? 1 : 0 )));  

//     BOOST_FOREACH(TAbsFitParticle* p, *m_particles) deletePart(p);
//     m_particles->clear();
//     
//     if (m_zllMassCons) { delete m_zllMassCons; m_zllMassCons = 0;}
//     if (m_zjjMassCons) { delete m_zjjMassCons; m_zjjMassCons = 0;}
//     if (m_pxCons) { delete m_pxCons; m_pxCons = 0;}
//     if (m_pyCons) { delete m_pyCons; m_pyCons = 0;}
    
  }

  void OSHKinFitter::deletePart(TAbsFitParticle* part) {

    if (part) {
      const TLorentzVector* vec = part->getIni4Vec();
      if (vec) delete vec;
      vec = 0;

      const TMatrixD* cov = part->getCovMatrix();
      if (cov) delete cov;
      cov = 0;

      delete part;
    }

    part = 0;
  }

  void OSHKinFitter::fit() {
    m_fitter.fit();

    if (m_fitter.getStatus() != 0) {
      ++m_nFailed;
      msg(MSG::WARNING) << "Fit Failed" << endreq; 
    }
  }

  void OSHKinFitter::setZll(HepLorentzVector& lep1, HepLorentzVector& lep2, int pdg, bool constrain) {

    float dE(0); 

    // Should we set mass to known mass?  Use gaus for mass if prob in BW?
    TLorentzVector* vec1 = hlvtotlv(lep1);

         if (abs(pdg)==11) dE = 0.1/sqrt(lep1.e()/GeV);
    else if (abs(pdg)==13) dE = (0.10/1000.)*(lep1.et()/GeV);
    TMatrixD* cov1 = energyCovMatrix(dE);    

    m_zllLep1 = new TFitParticleRelE( "ZllLep1", "ZllLep1", vec1, cov1);
    m_fitter.addMeasParticle(m_zllLep1);
    m_particles->push_back(m_zllLep1);

    TLorentzVector* vec2 = hlvtotlv(lep2);

    if (abs(pdg)==11) dE = 0.1/sqrt(lep2.e()/GeV);
    else if (abs(pdg)==13) dE = (0.10/1000.)*(lep2.et()/GeV);
    TMatrixD* cov2 = energyCovMatrix(dE);

    m_zllLep2 = new TFitParticleRelE( "ZllLep2", "ZllLep2", vec2, cov2);
    m_fitter.addMeasParticle(m_zllLep2);
    m_particles->push_back(m_zllLep2);

    if (constrain) {
      m_zllMassCons = new TFitConstraintMBW( "MassConstraint Zll", "Mass-Constraint Zll", 0, 0, Zmass, Zwidth);
      m_zllMassCons->addParticles1(m_zllLep1, m_zllLep2);
      m_fitter.addConstraint(m_zllMassCons);
    }
  }

  void OSHKinFitter::setZjj(HepLorentzVector& jet1, HepLorentzVector& jet2, bool constrain) {

    TLorentzVector* vec1 = hlvtotlv(jet1);
    // Should we set mass to known mass?
    TMatrixD* cov1 = energyCovMatrix(jet1.e() ? 0.8/sqrt(jet1.e()/GeV) : 0.8);

    m_zjjJet1 = new TFitParticleRelE( "ZjjJet1", "ZjjJet1", vec1, cov1);
    m_fitter.addMeasParticle(m_zjjJet1);
    m_particles->push_back(m_zjjJet1);

    TLorentzVector* vec2 = hlvtotlv(jet2);
    TMatrixD* cov2 = energyCovMatrix(jet2.e() ? 0.8/sqrt(jet2.e()/GeV) : 0.8);

    m_zjjJet2 = new TFitParticleRelE( "ZjjJet2", "ZjjJet2", vec2, cov2);
    m_fitter.addMeasParticle(m_zjjJet2);
    m_particles->push_back(m_zjjJet2);

    if (constrain) {
      m_zjjMassCons = new TFitConstraintMBW( "MassConstraint Zjj", "Mass-Constraint Zjj", 0, 0, Zmass, Zwidth);
      m_zjjMassCons->addParticles1(m_zjjJet1,m_zjjJet2 );
      m_fitter.addConstraint(m_zjjMassCons);
    }
  }

  void OSHKinFitter::setMETHZZllqq(std::vector<HepLorentzVector>& otherJets, float etMissX, float etMissY, bool constrain) {

    // Loop over other jets
    m_otherJets->reserve(otherJets.size());

    int ijet(0);
    BOOST_FOREACH (HepLorentzVector j, otherJets) {
      TLorentzVector* vec = hlvtotlv(j);
      TMatrixD* cov = energyCovMatrix(j.e() ? 0.8/sqrt(j.e()/GeV) : 0.8);
      std::stringstream name; name << "OtherJet" << ijet++;
      
      TFitParticleRelE* jet = new TFitParticleRelE(name.str().c_str(), name.str().c_str(), vec, cov);				      
    
      m_otherJets->push_back(jet);
      m_particles->push_back(jet);
      m_fitter.addMeasParticle(jet);
    }
    
    // Use MET to account for non-particle energy 
    // (must inlcude muons, should add the option of not doing so)
    TVector3 metVec(etMissX, etMissY, 0);
    //if (msgLvl(MSG::DEBUG)) {
    //msg(MSG::DEBUG) << "Initial MET" << endreq;
    std::cout << "Initial MET" << std::endl;
    metVec.Print();
    //}
    BOOST_FOREACH (TAbsFitParticle* p , *m_particles) {
      metVec -= TVector3(p->getIni4Vec()->Px(),p->getIni4Vec()->Py(), 0.);
    }
    //if (msgLvl(MSG::DEBUG)) {
    //msg(MSG::DEBUG) << "Remaining MET" << endreq;
    std::cout << "Remaining MET" << std::endl;  
    metVec.Print();
    //}
    TLorentzVector* met4Vec = new TLorentzVector(metVec,metVec.Mag());
    TMatrixD* metCov = energyCovMatrix(0.8/sqrt(met4Vec->E()/GeV));
    TFitParticleRelE* met = new TFitParticleRelE("MET_FitParticle","MET_FitParticle", met4Vec, metCov);
    m_particles->push_back(met);
    m_fitter.addMeasParticle(met);

    if (constrain) {
      m_pxCons =  new TFitConstraintEp("PXBalance","PXBalance",m_particles,TFitConstraintEp::pX, 0.);
      m_pyCons =  new TFitConstraintEp("PYBalance","PYBalance",m_particles,TFitConstraintEp::pY, 0.);

      m_fitter.addConstraint(m_pxCons);
      m_fitter.addConstraint(m_pyCons);
    }

  }


  std::vector<HepLorentzVector*>*  OSHKinFitter::getOtherJets() {
    std::vector<HepLorentzVector*>* jets = new std::vector<HepLorentzVector*>();
    jets->reserve(m_otherJets->size());

    BOOST_FOREACH(TAbsFitParticle* p , *m_otherJets) {
      jets->push_back(tlvtohlv(*const_cast<TLorentzVector*>(p->getCurr4Vec())));
    }
    return jets;
  }

  TMatrixD* OSHKinFitter::energyCovMatrix(double dE) {
    TMatrixD* cov = new TMatrixD(1,1);
    double dE2(dE*dE); 
    std::cout << "dE = " << dE << std::endl;
    (*cov)(0,0) = dE2;
    return cov;
  }


  TLorentzVector* OSHKinFitter::hlvtotlv(HepLorentzVector& hlv) {
    TLorentzVector* tlv = new TLorentzVector(hlv.px(),hlv.py(),hlv.pz(),hlv.e());
    return tlv;
  }

  HepLorentzVector* OSHKinFitter::tlvtohlv(TLorentzVector& tlv) {
    HepLorentzVector* hlv = new HepLorentzVector(tlv.Px(), tlv.Py(), tlv.Pz(), tlv.E());
    return hlv;
  }

  void OSHKinFitter::print() {

    std::cout << "==================== Results of Constraint Fit ===================" << std::endl;
    std::cout << "Status     = " << m_fitter.getStatus() << std::endl;
    std::cout << "Chi2       = " << m_fitter.getS() << " for NDF = " << m_fitter.getNDF() << std::endl;
    std::cout << "Constraint = " << m_fitter.getF() << std::endl;

    if (m_zllMassCons) {
      m_zllMassCons->print();
      m_zllLep1->getIni4Vec()->Print();
      m_zllLep1->getCurr4Vec()->Print();
      m_zllLep2->getIni4Vec()->Print();
      m_zllLep2->getCurr4Vec()->Print();
    }

    std::cout << "" << std::endl;

    if (m_zjjMassCons) {
      m_zjjMassCons->print();
      m_zjjJet1->getIni4Vec()->Print();
      m_zjjJet1->getCurr4Vec()->Print();
      m_zjjJet2->getIni4Vec()->Print();
      m_zjjJet2->getCurr4Vec()->Print();
    }

    std::cout << "" << std::endl;

    if (m_pxCons && m_pyCons) {
      m_pxCons->print();
      m_pyCons->print();      
    }

    m_fitter.print();

    std::cout << "==================================================================" << std::endl;
  }
  
}