// Michel Lefebvre <lefebvre@uvic.ca>
// version 1.00
#ifndef TOOLKINEFIT
#define TOOLKINEFIT

#include "TLorentzVector.h"
#include "TMinuit.h"
#include <vector>
#include <iostream>
#include <iomanip>

// global counters
// set in constructor
int gl_n = 0;    // number of calls to gl_fcn
int gl_nMax = 0; // max number of calls for which there is printing

// global pointer to data jets
// assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar
std::vector<TLorentzVector>* gl_jets = new std::vector<TLorentzVector>();

// global fit jets
// assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar
std::vector<TLorentzVector>* gl_jetsFit = new std::vector<TLorentzVector>();

// global function to obtain the jet parameter errors
void gl_errors(std::vector<TLorentzVector>& jets, std::vector<double>& eSigmas) {
  // for now use simple function of jet energy only
  eSigmas.clear();
  for (int i = 0; i < (int)jets.size(); ++i) {
    // use resolution = 100%/sqrt(E(GeV))
    double sigma = (jets[i].E() > 0.) ? sqrt(jets[i].E()) : 0.;
    eSigmas.push_back(sigma);
  }
}
 
// global function to be minimized
void gl_fcn(int &npar, double *gin, double &f, double *par, int iflag) {

  // setup a few constants
  const double MW = 80.40; // PDG W mass
  const double GW = 2.14;  // PDG W width
  //  const double GT = 1.5;   // estimated top width
  const double GT = 0.1;   // estimated top width for delta generation!
  const double GTroot2 = GT*sqrt(2.);

  // get the data
  // assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar
  std::vector<TLorentzVector> jets = *gl_jets;  // gl_jets is global
  
  // select case using iflag
  if (iflag == 1) {
    // initialization
  } else if (iflag == 2) {
    // compute derivatives and store them in gin
    gin[0] = (double)npar; // to quiet the compiler
  }
  
  // compute function
  // set fit jets vectors
  std::vector<TLorentzVector> jetsFit;
  jetsFit.clear();
  // assume jet mass is zero
  // energy is only parameter for now
  for (int i = 0; i < (int)jets.size(); ++i) {
    double s = par[i]/jets[i].E();
    TLorentzVector p(s*jets[i].Px(), s*jets[i].Py(), s*jets[i].Pz(), par[i]);
    jetsFit.push_back(p);
  }
  
  // mass constraints
  // beware: If M() is negative then -sqrt(-M()) is returned by TLorentzVector
  // compute fit W masses
  double mw0Fit = (jetsFit[0] + jetsFit[1]).M();
  double mw1Fit = (jetsFit[2] + jetsFit[3]).M();
  // compute difference between fit top masses
  double mt0Fit = (jetsFit[0] + jetsFit[1] + jetsFit[4]).M();
  double mt1Fit = (jetsFit[2] + jetsFit[3] + jetsFit[5]).M();
  double dmtFit = mt0Fit - mt1Fit;
  
  // get the error on the jet energy
  std::vector<double> eSigmas;
  gl_errors(jets, eSigmas);
  
  // compute chi2
  f = 0.;
  for (int i = 0; i < (int)jets.size(); ++i) {
    double d = (eSigmas[i] > 0.) ? (jets[i].E() - jetsFit[i].E()) / eSigmas[i] : 0.;
    f += d*d;
  }
  // add simple gaussian W mass constraints
  f += pow((mw0Fit - MW) / GW, 2);
  f += pow((mw1Fit - MW) / GW, 2);
  // add constraint on the difference of the top masses
  f += pow(dmtFit/GTroot2, 2);

  if (iflag == 3) {
    // when the fit is finished
    // store fit jet vectors globally for convenience
    gl_jetsFit->clear();
    for (int i = 0; i < (int)jetsFit.size(); ++i) gl_jetsFit->push_back(jetsFit[i]);
    // std::cout << "gl_fcn: fit jets stored" << std::endl;
  }

  // extra printout
  if (gl_n < gl_nMax) {
    const int w = 10;
    if (gl_n%10 == 0) {
      std::cout << std::setw(w) << "npar"
		<< std::setw(w) << "iflag"
		<< std::setw(w) << "mw0Fit" 
		<< std::setw(w) << "mw1Fit"
		<< std::setw(w) << "mt0Fit" 
		<< std::setw(w) << "mt1Fit" 
		<< std::setw(w) << "f"  << std::endl;
    }     
    std::cout << std::setw(w) << npar
	      << std::setw(w) << iflag
	      << std::setw(w) << mw0Fit
	      << std::setw(w) << mw1Fit
	      << std::setw(w) << mt0Fit
	      << std::setw(w) << mt1Fit
	      << std::setw(w) << f << std::endl;
  }

  // count number of calls.  gl_n is global
  ++gl_n;
}


class toolKineFit {

 public:
  // ***** public methods
  toolKineFit();
  ~toolKineFit();

  // set the data jet four vectors
  // assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar
  void setJets(std::vector<TLorentzVector>& jets);

  // set print flag
  // level:  MINUIT SET PRIntout
  // ncalls: extra printout for first ncalls
  void setPrintLevel(int level = 0, int ncalls = 0);

  // perform kinematic fit
  bool fit();

  // fit is successful if full accurate covariance matrix obtained
  bool isValid();

  // get the fit jet four vectors
  // assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar
  std::vector<TLorentzVector> getJetsFit();

  // get the minimum value of the function to minimize
  double getFcnMin();

 private:
  // ***** private members
};

#endif