// Michel Lefebvre <lefebvre@uvic.ca>
#include "toolKineFit.h"
#include "utilIndex.h"
#include "TROOT.h"
#include <iostream>
#include <iomanip>
#include <vector>

void kinePrint(std::vector<TLorentzVector> partons, 
	       std::vector<TLorentzVector> jets, 
	       std::vector<TLorentzVector> jetsFit, 
	       std::vector<double> pull) {

  // print comparisons and the energy pull (the first 6 elements of pull)
  int w = 15;
  std::cout << std::setw(w) << " " 
	    << std::setw(w) << "px" << std::setw(w) << "py" << std::setw(w) << "pz" 
	    << std::setw(w) << "e"  << std::setw(w) << "m"  << std::setw(w) << "Epull" << std::endl;
  for (int j = 0; j < (int)partons.size(); ++j) {
    // print the six partons
    std::cout << std::setw(w) << "truth"
	      << std::setw(w) << partons[j].Px()
	      << std::setw(w) << partons[j].Py()
	      << std::setw(w) << partons[j].Pz()
	      << std::setw(w) << partons[j].E() 
	      << std::setw(w) << partons[j].M() 
	      << std::endl;
    // print the jets in hypothesis order (same order as partons)
    std::cout << std::setw(w) << "hypo"
	      << std::setw(w) << jets[j].Px()
	      << std::setw(w) << jets[j].Py()
	      << std::setw(w) << jets[j].Pz()
	      << std::setw(w) << jets[j].E() 
	      << std::setw(w) << jets[j].M() 
	      << std::endl;
    // print the fit jets (same order as partons)
    std::cout << std::setw(w) << "fit"
	      << std::setw(w) << jetsFit[j].Px()
	      << std::setw(w) << jetsFit[j].Py()
	      << std::setw(w) << jetsFit[j].Pz()
	      << std::setw(w) << jetsFit[j].E() 
	      << std::setw(w) << jetsFit[j].M() 
	      << std::setw(w) << pull[j]
	      << std::endl;
    std::cout << std::endl;
  }
  
  // print a few masses
  std::cout << std::setw(w) << " " 
	    << std::setw(w) << "mwplus" << std::setw(w) << "mwminus" 
	    << std::setw(w) << "mtop"   << std::setw(w) << "mtopbar"  
	    << std::endl;
  // truth
  std::cout << std::setw(w) << "truth"
	    << std::setw(w) << (partons[gl_Iu   ] + partons[gl_Idbar]).M()
	    << std::setw(w) << (partons[gl_Iubar] + partons[gl_Id   ]).M()
	    << std::setw(w) << (partons[gl_Iu   ] + partons[gl_Idbar] + partons[gl_Ib   ]).M()
	    << std::setw(w) << (partons[gl_Iubar] + partons[gl_Id   ] + partons[gl_Ibbar]).M()
	    << std::endl;
  // hypothesis
  std::cout << std::setw(w) << "hypo"
	    << std::setw(w) << (jets[gl_Iu   ] + jets[gl_Idbar]).M()
	    << std::setw(w) << (jets[gl_Iubar] + jets[gl_Id   ]).M()
	    << std::setw(w) << (jets[gl_Iu   ] + jets[gl_Idbar] + jets[gl_Ib   ]).M()
	    << std::setw(w) << (jets[gl_Iubar] + jets[gl_Id   ] + jets[gl_Ibbar]).M()
	    << std::endl;
  // fit
  std::cout << std::setw(w) << "fit"
	    << std::setw(w) << (jetsFit[gl_Iu   ] + jetsFit[gl_Idbar]).M()
	    << std::setw(w) << (jetsFit[gl_Iubar] + jetsFit[gl_Id   ]).M()
	    << std::setw(w) << (jetsFit[gl_Iu   ] + jetsFit[gl_Idbar] + jetsFit[gl_Ib   ]).M()
	    << std::setw(w) << (jetsFit[gl_Iubar] + jetsFit[gl_Id   ] + jetsFit[gl_Ibbar]).M()
	    << std::endl;
}


void runKineFit() {
  std::cout << "entering kineFit" << std::endl;
  
  // enter an event by hand (px, py, pz, e)
  // see utilIndex.h for index values
  std::vector<TLorentzVector> jets;
  jets.resize(gl_Npartons);
  jets[gl_Iu   ] = TLorentzVector( 83.6813,   28.7426,   318.708,   330.762);  // u
  jets[gl_Idbar] = TLorentzVector(-5.51842,    34.728,   357.962,   359.685);  // dbar
  jets[gl_Iubar] = TLorentzVector(  23.952,   -68.826,   63.5647,   96.7015);  // ubar
  jets[gl_Id   ] = TLorentzVector( 78.2212,  -45.8589,    183.21,   204.419);  // d
  jets[gl_Ib   ] = TLorentzVector(-25.0896,   75.4045,   176.077,    193.18);  // b
  jets[gl_Ibbar] = TLorentzVector(-37.3136,   9.49893,   39.1253,   54.8938);  // bbar

  // now enter the corresponding parton truth for later comparison
  std::vector<TLorentzVector> partons;
  partons.resize(gl_Npartons);
  partons[gl_Iu   ] = TLorentzVector( 76.7362,   26.7729,    281.28,   292.786);  // u   
  partons[gl_Idbar] = TLorentzVector(-4.47932,   33.7419,   357.843,   359.459);  // dbar
  partons[gl_Iubar] = TLorentzVector( 22.8382,  -73.7662,   72.5491,   105.955);  // ubar
  partons[gl_Id   ] = TLorentzVector( 74.8944,  -42.0846,   175.085,   195.026);  // d   
  partons[gl_Ib   ] = TLorentzVector(-23.4837,    77.478,   174.558,   192.484);  // b   
  partons[gl_Ibbar] = TLorentzVector(-39.3133,   8.36382,   47.7447,   62.6102);  // bbar

  // prepare fitter
  toolKineFit kineFit;
  std::vector<TLorentzVector> jetsFit;
  std::vector<double> pull;

  // ******

  // try type 0 fit
  kineFit.setFitType(0);     // type 0, only vary energy of jets (6 parameters)
  kineFit.setPrintLevel(0);  // you can try setPrintLevel(1, 100) for more output!
  kineFit.setJets(jets);
  std::cout << std::endl;
  std::cout << "***** type 0 fit ******" << std::endl;
  // perform fit
  kineFit.fit();
  // look at results
  // is the fit valid?
  if (kineFit.isValid()) {
    std::cout << "Fit is valid" << std::endl;
  } else {
    std::cout << "Fit is NOT valid!!" << std::endl;
  }
  // get the minimum of the fitting fucntion (here it's a chi2)
  std::cout << "chi2 = " << kineFit.getFcnMin() << std::endl;
  // get the jets after the constrained fit and the pull
  jetsFit = kineFit.getJetsFit();
  pull = kineFit.getFitPull();   // here pull has 6 elements (energy pull)
  // print comparisons and energy pull
  kinePrint(partons, jets, jetsFit, pull);

  // *****
  // try type 1 fit
  kineFit.setFitType(1);     // type 1, vary energy, eta, phi of jets (18 parameters)
  kineFit.setPrintLevel(0);  // you can try setPrintLevel(1, 100) for more output!
  kineFit.setJets(jets);
  std::cout << std::endl;
  std::cout << "***** type 1 fit ******" << std::endl;
  // perform fit
  kineFit.fit();
  // look at results
  // is the fit valid?
  if (kineFit.isValid()) {
    std::cout << "Fit is valid" << std::endl;
  } else {
    std::cout << "Fit is NOT valid!!" << std::endl;
  }
  // get the minimum of the fitting fucntion (here it's a chi2)
  std::cout << "chi2 = " << kineFit.getFcnMin() << std::endl;
  // get the jets after the constrained fit
  jetsFit = kineFit.getJetsFit();
  pull = kineFit.getFitPull();   // here pull has 18 elements (energy, eta, phi pull)
  // print comparisons
  kinePrint(partons, jets, jetsFit, pull);

}