// Michel Lefebvre <lefebvre@uvic.ca>
#include "toolKineFit.h"
#include "TROOT.h"
#include "TString.h"

// *****
toolKineFit::toolKineFit(int type) {
  // set global constants
  gl_n = 0;
  gl_nMax = 0;
  if (gl_jets) delete gl_jets;
  gl_jets = new std::vector<TLorentzVector>();
  if (gl_jetsFit) delete gl_jetsFit;
  gl_jetsFit = new std::vector<TLorentzVector>();

  // set fit type
  this->setFitType(type);
}

// *****
toolKineFit::~toolKineFit() {
  delete gl_jets;
  delete gl_jetsFit;
}

// *****
void
toolKineFit::setFitType(int type) {
  m_fitType = type;

  if (m_fitType == 1) {
    // type 1
    // simple chi2 fit, varying energy, eta and phi of jets
    // initialize minuit, which is global
    if (gMinuit) delete gMinuit;
    gMinuit = new TMinuit(18); // 18 parameters: 0-5 energies, 6-11 etas, 12-17 phis
    gMinuit->SetFCN(gl_fcn1);  // gl_fcn1 is global

  } else {
    if (m_fitType != 0 ) {
      std::cout << "toolKineFit extentiated with bad type: " << m_fitType << std::endl;
      std::cout << "assume type 0 and proceed" << std::endl;
      m_fitType = 0;
    }
    // type 0
    // simple chi2 fit, varying only energy of jets
    // initialize minuit, which is global
    if (gMinuit) delete gMinuit;
    gMinuit = new TMinuit(6); // 6 parameters: 0-5 energies
    gMinuit->SetFCN(gl_fcn0); // gl_fcn0 is global
  }
}
  
// *****
void
toolKineFit::setJets(std::vector<TLorentzVector>& jets) {
  // gl_jets is global
  gl_jets->clear();
  for (int j = 0; j < (int)jets.size(); ++j) gl_jets->push_back(jets[j]);
}

// *****
void
toolKineFit::setPrintLevel(int level, int ncalls) {
  gl_nMax = ncalls;
  int ierflg = 0;
  double arglist[1];
  arglist[0] = level;
  gMinuit->mnexcm("SET PRI", arglist, 1, ierflg);
}

// *****
bool
toolKineFit::fit() {
  // get the data
  std::vector<TLorentzVector> jets = *gl_jets;  // gl_jets is global

  // number of jets
  int nJet = (int)jets.size();

  // get the error on the jet energy, eta and phi
  std::vector<double> eSigmas, etaSigmas, phiSigmas;
  gl_errors(jets, eSigmas, etaSigmas, phiSigmas);

  // define starting values and step sizes for parameters
  int ierflg = 0;
  double start, step, minVal, maxVal;
  for (int i = 0; i < nJet; ++i) {
    if (m_fitType == 0 || m_fitType == 1) {
      // energy
      start = jets[i].E(); // guess initial value
      step = eSigmas[i]; // guess initial error
      // if minVal = maxVal = 0 then the parameter is unbounded
      minVal = 0.;
      maxVal = 1000000.;
      TString t = "e"; t += i;
      ierflg = 0;
      gMinuit->mnparm(i, t, start, step, minVal, maxVal, ierflg);
      // std::cout << "mnparm " << i << " ierflg = " << ierflg << std::endl;
    }
    if (m_fitType == 1) {
      // eta
      start = jets[i].Eta(); // guess initial value
      step = etaSigmas[i]; // guess initial error
      // if minVal = maxVal = 0 then the parameter is unbounded
      minVal = 0.;
      maxVal = 0.;
      TString t = "eta"; t += i;
      ierflg = 0;
      gMinuit->mnparm(i+nJet, t, start, step, minVal, maxVal, ierflg);
      // std::cout << "mnparm " << i << " ierflg = " << ierflg << std::endl;

      // phi
      start = jets[i].Phi(); // guess initial value
      step = phiSigmas[i]; // guess initial error
      // if minVal = maxVal = 0 then the parameter is unbounded
      minVal = -3.1416;
      maxVal = 3.1416;
      t = "phi"; t += i;
      ierflg = 0;
      gMinuit->mnparm(i+2*nJet, t, start, step, minVal, maxVal, ierflg);
      // std::cout << "mnparm " << i << " ierflg = " << ierflg << std::endl;
    }
  }

  // chi2 "up" parameter for error definition
  gMinuit->SetErrorDef(1);
  // start minimization
  gMinuit->Migrad();
  //
  return true;
}

// *****
bool
toolKineFit::isValid() {
  double fmin, fedm, errdef;
  int nvpar, nparx, icstat;
  gMinuit->mnstat(fmin, fedm, errdef, nvpar, nparx, icstat);
  // std::cout << "icstat = " << icstat << std::endl;
  return icstat==3;
}
  

// *****
std::vector<TLorentzVector>
toolKineFit::getJetsFit() {
  // call function with termination flag to fill vector of jetFits
  int ierflg = 0;
  double arglist[1] = {3};
  gMinuit->mnexcm("CAL1", arglist, 1, ierflg);
  return *gl_jetsFit;
}

// *****
double
toolKineFit::getFcnMin() {
  double fmin, fedm, errdef;
  int nvpar, nparx, icstat;
  gMinuit->mnstat(fmin, fedm, errdef, nvpar, nparx, icstat);
  return fmin;
}

// *****
std::vector<double>
toolKineFit::getFitPull() {
  // get the measured jet four vectors
  std::vector<TLorentzVector> jets = *gl_jets;  // gl_jets is global
  
  // number of jets
  int nJet = (int)jets.size();

  // get the error on the jet energy, eta and phi
  std::vector<double> eSigmas, etaSigmas, phiSigmas;
  gl_errors(jets, eSigmas, etaSigmas, phiSigmas);  // gl_error is global

  // get the fit jet four vectors
  std::vector<TLorentzVector> jetsFit = this->getJetsFit();

  std::vector<double> pull;
  // fit types 0 and 1
  if (m_fitType == 0 || m_fitType == 1) {
    // energies
    for (int i = 0; i < nJet; ++i) {
      double d = (eSigmas[i] > 0.) ? (jets[i].E() - jetsFit[i].E()) / eSigmas[i] : 0.;
      pull.push_back(d);
    }
  }
  // fit type 1
  if (m_fitType == 1) {
    for (int i = 0; i < nJet; ++i) {
      // etas
      double d = (etaSigmas[i] > 0.) ? (jets[i].Eta() - jetsFit[i].Eta()) / etaSigmas[i] : 0.;
      pull.push_back(d);
   }
    for (int i = 0; i < nJet; ++i) {
      // phis
      double d = (phiSigmas[i] > 0.) ? (jets[i].Phi() - jetsFit[i].Phi()) / phiSigmas[i] : 0.;
      pull.push_back(d);
    }
  }

  return pull;
}