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

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

// assumed hypothesis order set in utilIndex.h
#include "utilIndex.h"

// setup a few global constants
const double gl_MW = 80.40; // PDG W mass
const double gl_GW = 2.14;  // PDG W width
//  const double gl_GT = 1.5;   // estimated top width
const double gl_GT = 0.1;   // estimated top width for delta generation!
const double gl_GTroot2 = gl_GT * sqrt(2.);

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

// global pointer to data jets
std::vector<TLorentzVector>* gl_jets = new std::vector<TLorentzVector>();

// global fit jets
std::vector<TLorentzVector>* gl_jetsFit = new std::vector<TLorentzVector>();

// global function to obtain the jet components errors
void gl_errors(std::vector<TLorentzVector>& jets, 
	       std::vector<double>& eSigmas, std::vector<double>& etaSigmas, std::vector<double>& phiSigmas) {
  // for now use simple function of jet energy
  // for now, use fixed error in eta and phi
  eSigmas.clear();
  etaSigmas.clear();
  phiSigmas.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);
    // for now use 0.06 as rms in eta and phi
    etaSigmas.push_back(0.06);
    phiSigmas.push_back(0.06);
  }
}
 
// fit type 0: global function to be minimized
void gl_fcn0(int &npar, double *gin, double &f, double *par, int iflag) {

  // get the data
  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
  // vary only energy
  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[gl_Iu]    + jetsFit[gl_Idbar]).M();
  double mw1Fit = (jetsFit[gl_Iubar] + jetsFit[gl_Id]).M();
  // compute difference between fit top masses
  double mt0Fit = (jetsFit[gl_Iu]    + jetsFit[gl_Idbar] + jetsFit[gl_Ib]).M();
  double mt1Fit = (jetsFit[gl_Iubar] + jetsFit[gl_Id]    + jetsFit[gl_Ibbar]).M();
  double dmtFit = mt0Fit - mt1Fit;
  
  // get the error on the jet components
  std::vector<double> eSigmas, etaSigmas, phiSigmas;
  gl_errors(jets, eSigmas, etaSigmas, phiSigmas);
  
  // 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 - gl_MW) / gl_GW, 2);
  f += pow((mw1Fit - gl_MW) / gl_GW, 2);
  // add constraint on the difference of the top masses
  f += pow(dmtFit/gl_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_fcn0: 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;
}

void gl_fcn1(int &npar, double *gin, double &f, double *par, int iflag) {

  // get the data
  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
  // vary energy, eta, phi of jets
  for (int i = 0; i < (int)jets.size(); ++i) {
    double e = par[i];
    double eta = par[i+6];
    double phi = par[i+12];
    double pt = e/cosh(eta);  // cosh(eta) = 1/sin(theta) >=1 for 0 <= theta <= pi
    double px = pt * cos(phi);
    double py = pt * sin(phi);
    double pz = e * tanh(eta); // tanh(eta) = cos(theta) 
    TLorentzVector p(px, py, pz, e);
    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[gl_Iu]    + jetsFit[gl_Idbar]).M();
  double mw1Fit = (jetsFit[gl_Iubar] + jetsFit[gl_Id]).M();
  // compute difference between fit top masses
  double mt0Fit = (jetsFit[gl_Iu]    + jetsFit[gl_Idbar] + jetsFit[gl_Ib]).M();
  double mt1Fit = (jetsFit[gl_Iubar] + jetsFit[gl_Id]    + jetsFit[gl_Ibbar]).M();
  double dmtFit = mt0Fit - mt1Fit;
  
  // get the error on the jet components
  std::vector<double> eSigmas, etaSigmas, phiSigmas;
  gl_errors(jets, eSigmas, etaSigmas, phiSigmas);
  
  // 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;
    d = (etaSigmas[i] > 0.) ? (jets[i].Eta() - jetsFit[i].Eta()) / etaSigmas[i] : 0.;
    f += d*d;
    d = (phiSigmas[i] > 0.) ? (jets[i].Phi() - jetsFit[i].Phi()) / phiSigmas[i] : 0.;
    f += d*d;
  }
  // add simple gaussian W mass constraints
  f += pow((mw0Fit - gl_MW) / gl_GW, 2);
  f += pow((mw1Fit - gl_MW) / gl_GW, 2);
  // add constraint on the difference of the top masses
  f += pow(dmtFit/gl_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_fcn1: 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(int type = 0);
  ~toolKineFit();

  // set the type of fit
  void setFitType(int type = 0);

  // set the data jet four vectors
  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
  std::vector<TLorentzVector> getJetsFit();

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

 private:
  // ***** private members

  // fit type
  int m_fitType;
};

#endif