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

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

  // print comparisons
  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::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::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[0] + partons[1]).M()
	    << std::setw(w) << (partons[2] + partons[3]).M()
	    << std::setw(w) << (partons[0] + partons[1] + partons[4]).M()
	    << std::setw(w) << (partons[2] + partons[3] + partons[5]).M()
	    << std::endl;
  // hypothesis
  std::cout << std::setw(w) << "hypo"
	    << std::setw(w) << (jets[0] + jets[1]).M()
	    << std::setw(w) << (jets[2] + jets[3]).M()
	    << std::setw(w) << (jets[0] + jets[1] + jets[4]).M()
	    << std::setw(w) << (jets[2] + jets[3] + jets[5]).M()
	    << std::endl;
  // fit
  std::cout << std::setw(w) << "fit"
	    << std::setw(w) << (jetsFit[0] + jetsFit[1]).M()
	    << std::setw(w) << (jetsFit[2] + jetsFit[3]).M()
	    << std::setw(w) << (jetsFit[0] + jetsFit[1] + jetsFit[4]).M()
	    << std::setw(w) << (jetsFit[2] + jetsFit[3] + jetsFit[5]).M()
	    << std::endl;
}


void runKineFit() {
  std::cout << "entering kineFit with fit" << std::endl;
  
  // enter an event by hand (px, py, pz, e)
  // assumed hypothesis order: 0=u  1=dbar  2=ubar  3=d  4=b  5=bbar  -> set in utilIndex.h
  std::vector<TLorentzVector> jets;                                         
  jets.push_back(TLorentzVector( 83.6813,   28.7426,   318.708,   330.762));
  jets.push_back(TLorentzVector(-5.51842,    34.728,   357.962,   359.685));
  jets.push_back(TLorentzVector(  23.952,   -68.826,   63.5647,   96.7015));
  jets.push_back(TLorentzVector( 78.2212,  -45.8589,    183.21,   204.419));
  jets.push_back(TLorentzVector(-25.0896,   75.4045,   176.077,    193.18));
  jets.push_back(TLorentzVector(-37.3136,   9.49893,   39.1253,   54.8938));

  // now enter the corresponding parton truth for later comparison
  std::vector<TLorentzVector> partons;
  partons.push_back(TLorentzVector( 76.7362,   26.7729,    281.28,   292.786));
  partons.push_back(TLorentzVector(-4.47932,   33.7419,   357.843,   359.459));
  partons.push_back(TLorentzVector( 22.8382,  -73.7662,   72.5491,   105.955));
  partons.push_back(TLorentzVector( 74.8944,  -42.0846,   175.085,   195.026));
  partons.push_back(TLorentzVector(-23.4837,    77.478,   174.558,   192.484));
  partons.push_back(TLorentzVector(-39.3133,   8.36382,   47.7447,   62.6102));

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

  // ******

  // 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
  jetsFit = kineFit.getJetsFit();
  // print comparisons
  kinePrint(partons, jets, jetsFit);

  // *****
  // 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();
  // print comparisons
  kinePrint(partons, jets, jetsFit);

}