#include <iostream>
#include <vector>
#include <map>

#include "TTree.h"
#include "TFile.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TF1.h"
#include "TRandom3.h"
#include "TLeaf.h"

#include "tmb_tree/TMBLorentzVector.hpp"
#include "top_cafe/TopCreateObjects.hpp"
#include "top_cafe/TopTopologicalVariables.hpp"

void    ProcessBar(int thisevt, int totalevt); 
int     CalculateNuPz(double w_mass, 
		      const TMBLorentzVector &lep, 
		      double MET_x, double MET_y, double &S1, double &S2);
double *ReconstructTop(const TMBLorentzVector &lep,
                       const TMBLorentzVector &jet,
                       double MET_x, double MET_y);
double *ReconstructTop(const TMBLorentzVector &lep,
		       std::vector<TMBLorentzVector> &jets,
		       double MET_x, double MET_y, int ijet);
TMBLorentzVector MakeTop(const TMBLorentzVector &lep,
                         const TMBLorentzVector &nu,
                         const TMBLorentzVector &jet);
TF1      * fun_sigma;
TRandom3 * rndm;
const double W_Mass = 80.43;
const bool   debug  = true;
double optvar[6];

int main(int argc, const char* argv[]) {

  //jet energy resolution sigma curve
  fun_sigma = new TF1("fun_sigma", "[0]+[1]*TMath::Exp(-[2]*x)", 0, 1000);
  fun_sigma->FixParameter(0, 1.63325e-01);
  fun_sigma->FixParameter(1, 2.27338e-01);
  fun_sigma->FixParameter(2, 2.11078e-02);
  //fun_sigma->Write();
  rndm = new TRandom3();

  //TFile * inputfile = new TFile("p17_CC_ttbar_Test_Topo.root", "READ");
  std::cout << "Input file name: " << argv[1] << std::endl;
  TFile * inputfile = new TFile(argv[1], "UPDATE");
  TTree * tree = (TTree*)inputfile->Get("TopologicalVariables");
  int Ntotal = tree->GetEntries();

  tree->SetBranchStatus("*",       0);
  tree->SetBranchStatus("Jet*",    1);
  tree->SetBranchStatus("Lepton*", 1);
  tree->SetBranchStatus("MET*",    1);
  tree->SetBranchStatus("NGoodJets",    1);
  tree->SetBranchStatus("EventNumber",    1);

  TString str_array[] = {"LeptonPx", "LeptonPy", "LeptonPz", "LeptonE", 
			 "METPx", "METPy"};
  int Nstr = sizeof(str_array)/sizeof(str_array[0]);
  std::map<TString, Double_t> _Doubles;
  for (int istr=0; istr<Nstr; istr++) 
    tree->SetBranchAddress(str_array[istr], &_Doubles[str_array[istr]]);
  int NGoodJets;
  tree->SetBranchAddress("NGoodJets", &NGoodJets);

  struct OptVars {
    double OrgNuPz_S1;
    double OrgNuPz_S2;   
    double OptJet1NuPz_S1;
    double OptJet1NuPz_S2;
    double OptJet1TopMass_S1;
    double OptJet1TopMass_S2;
    double OptJet1TM_Posterior_S1;
    double OptJet1TM_Posterior_S2;
    double OptJet2NuPz_S1;
    double OptJet2NuPz_S2;
    double OptJet2TopMass_S1;
    double OptJet2TopMass_S2;
    double OptJet2TM_Posterior_S1;
    double OptJet2TM_Posterior_S2;
    double OptJet3NuPz_S1;
    double OptJet3NuPz_S2;
    double OptJet3TopMass_S1;
    double OptJet3TopMass_S2;
    double OptJet3TM_Posterior_S1;
    double OptJet3TM_Posterior_S2;
    double OptJet4NuPz_S1;
    double OptJet4NuPz_S2;
    double OptJet4TopMass_S1;
    double OptJet4TopMass_S2;
    double OptJet4TM_Posterior_S1;
    double OptJet4TM_Posterior_S2;
  };

  OptVars myopt;
  TBranch * newbr = tree->Branch("OptVars", &myopt.OrgNuPz_S1, "OrgNuPz_S1/D:OrgNuPz_S2:OptJet1NuPz_S1:OptJet1NuPz_S2:OptJet1TopMass_S1:OptJet1TopMass_S2:OptJet1TM_Posterior_S1:OptJet1TM_Posterior_S2:OptJet2NuPz_S1:OptJet2NuPz_S2:OptJet2TopMass_S1:OptJet2TopMass_S2:OptJet2TM_Posterior_S1:OptJet2TM_Posterior_S2:OptJet3NuPz_S1:OptJet3NuPz_S2:OptJet3TopMass_S1:OptJet3TopMass_S2:OptJet3TM_Posterior_S1:OptJet3TM_Posterior_S2:OptJet4NuPz_S1:OptJet4NuPz_S2:OptJet4TopMass_S1:OptJet4TopMass_S2:OptJet4TM_Posterior_S1:OptJet4TM_Posterior_S2");

  for (int ievt=0; ievt<Ntotal; ievt++) {
  //for (int ievt=0; ievt<20; ievt++) {
    ProcessBar(ievt, Ntotal); //print status bar
    tree->GetEntry(ievt);
    
    TMBLorentzVector lep;
    lep.SetPxPyPzE(_Doubles["LeptonPx"], 
		   _Doubles["LeptonPy"], 
		   _Doubles["LeptonPz"], 
		   _Doubles["LeptonE"]);

    myopt.OrgNuPz_S1 = -999.;
    myopt.OrgNuPz_S2 = -999.;
    CalculateNuPz(W_Mass, lep, _Doubles["METPx"], _Doubles["METPy"], 
		  myopt.OrgNuPz_S1, myopt.OrgNuPz_S2);
    
    std::vector<TMBLorentzVector> vec_jets;
    for (int ijet=0; ijet<(NGoodJets>4?3:NGoodJets); ijet++) {
      TMBLorentzVector jet;
      double jetpx = tree->GetLeaf(Form("Jet%iPx", ijet+1))->GetValue();
      double jetpy = tree->GetLeaf(Form("Jet%iPy", ijet+1))->GetValue();
      double jetpz = tree->GetLeaf(Form("Jet%iPz", ijet+1))->GetValue(); 
      double jete  = tree->GetLeaf(Form("Jet%iE",  ijet+1))->GetValue();
      jet.SetPxPyPzE(jetpx, jetpy, jetpz, jete);
      vec_jets.push_back(jet);
    }

    //initialize the struct
    myopt.OptJet1NuPz_S1           = -999;
    myopt.OptJet1NuPz_S2           = -999;
    myopt.OptJet1TopMass_S1        = -999;
    myopt.OptJet1TopMass_S2        = -999;
    myopt.OptJet1TM_Posterior_S1   = -999;
    myopt.OptJet1TM_Posterior_S2   = -999;
    myopt.OptJet2NuPz_S1           = -999;
    myopt.OptJet2NuPz_S2           = -999;
    myopt.OptJet2TopMass_S1        = -999;
    myopt.OptJet2TopMass_S2        = -999;
    myopt.OptJet2TM_Posterior_S1   = -999;
    myopt.OptJet2TM_Posterior_S2   = -999;
    myopt.OptJet3NuPz_S1           = -999;
    myopt.OptJet3NuPz_S2           = -999;
    myopt.OptJet3TopMass_S1        = -999;
    myopt.OptJet3TopMass_S2        = -999;
    myopt.OptJet3TM_Posterior_S1   = -999;
    myopt.OptJet3TM_Posterior_S2   = -999;
    myopt.OptJet4NuPz_S1           = -999;
    myopt.OptJet4NuPz_S2           = -999;
    myopt.OptJet4TopMass_S1        = -999;
    myopt.OptJet4TopMass_S2        = -999;
    myopt.OptJet4TM_Posterior_S1   = -999;
    myopt.OptJet4TM_Posterior_S2   = -999;


    if (ievt>=20) {
      newbr->Fill();
      continue;
    }
    for (int jjet=0; jjet<vec_jets.size(); jjet++) {
      double * _optvar;
      _optvar = ReconstructTop(lep, vec_jets, _Doubles["METPx"], _Doubles["METPy"], jjet);
      
      if ( vec_jets.size()>=1 || jjet==0 ) {
	myopt.OptJet1NuPz_S1           = _optvar[0];
	myopt.OptJet1NuPz_S2           = _optvar[1];
	myopt.OptJet1TopMass_S1        = _optvar[2];
	myopt.OptJet1TopMass_S2        = _optvar[3];
	myopt.OptJet1TM_Posterior_S1   = _optvar[4];
	myopt.OptJet1TM_Posterior_S2   = _optvar[5];
      }
      if (vec_jets.size()>=2 || jjet==1 ) {
	myopt.OptJet2NuPz_S1           = _optvar[0];
	myopt.OptJet2NuPz_S2           = _optvar[1];
	myopt.OptJet2TopMass_S1        = _optvar[2];
	myopt.OptJet2TopMass_S2        = _optvar[3];
	myopt.OptJet2TM_Posterior_S1   = _optvar[4];
	myopt.OptJet2TM_Posterior_S2   = _optvar[5];
      }
      if (vec_jets.size()>=3 || jjet==2 ) {
	myopt.OptJet3NuPz_S1           = _optvar[0];
	myopt.OptJet3NuPz_S2           = _optvar[1];
	myopt.OptJet3TopMass_S1        = _optvar[2];
	myopt.OptJet3TopMass_S2        = _optvar[3];
	myopt.OptJet3TM_Posterior_S1   = _optvar[4];
	myopt.OptJet3TM_Posterior_S2   = _optvar[5];
      }
      if (vec_jets.size()>=4 || jjet==3 ) {
	myopt.OptJet4NuPz_S1           = _optvar[0];
	myopt.OptJet4NuPz_S2           = _optvar[1];
	myopt.OptJet4TopMass_S1        = _optvar[2];
	myopt.OptJet4TopMass_S2        = _optvar[3];
	myopt.OptJet4TM_Posterior_S1   = _optvar[4];
	myopt.OptJet4TM_Posterior_S2   = _optvar[5];	
      }
      if (debug) {
	printf("\nJet %i", jjet);
	for (int i=0; i<6; i++) {
	  printf("%8.3f  ", (float)_optvar[i]);
	}
      }
    }
    if (debug)
      printf("== Lepton Px: %7.3f\n", (float)_Doubles["LeptonPx"]);
    newbr->Fill();    
  }
  tree->Write("", TObject::kOverwrite);

}

void ProcessBar(int thisevt, int totalevt) {
  if ( thisevt%500 != 0 ) return ;
  int totalc = 30;
  int ic = (int)1.*totalc*thisevt/totalevt;
  std::cout << "Status:   [";
  for (int a=0; a<ic; a++)        std::cout << "=";
  for (int b=0; b<totalc-ic; b++) std::cout << " ";
  std::cout << "]";

  printf("%6.2f%%", 100.*thisevt/totalevt);

  //backspace
  for(int a = 0; a < totalc+19; a++) std::cout << "\b";

}

int CalculateNuPz(double w_mass, const TMBLorentzVector &lep, double MET_x, double MET_y, double &S1, double &S2) {

  top_cafe::TopCreateObjects temp;
  if(!temp.isValid(lep)) return -1;
  // variables needed for the W reconstruction:
  double A, B, C1, C, k, S_denom, mwplz, scf, Mt, charge;

  //calculate the Pz of the neutrino
  TMBLorentzVector nu;
  nu.SetXYZM(MET_x,MET_y,0.0,0.0);

  std::vector<TMBLorentzVector> lepMETxMETy;
  lepMETxMETy.push_back(lep);
  lepMETxMETy.push_back(nu);
  temp.SetInputLeptonMET_W(lepMETxMETy);

  top_cafe::TopTopologicalVariables leptonMETxy(lepMETxMETy);
  Mt = leptonMETxy.TransverseMass();

  //     if () return -1;
  scf   = 1.0;

  if (Mt < w_mass)
    A = w_mass*w_mass / 2.;
  else {           // assume Mt=w_mass, rescale MET accordingly
    A = Mt*Mt / 2.;
    k = nu.Pt()*lep.Pt() - nu.Px()*lep.Px() - nu.Py()*lep.Py();
    k = (k == 0 ? 0.00001 : k);
    scf = 0.5 * w_mass*w_mass / k;
    nu *= scf;
  }
  int Nsolution = 0; //how many solutions
  // if the lepton goes along the beam direction (Px==Py==0),
  // then the algorithm below doesn't work, so do the calculation here.
  //S_denom = pow(l_pz, 2) - pow(l_e, 2);
  S_denom = lep.Pz()*lep.Pz() - lep.E()*lep.E();
  if(S_denom == 0.) {
    // in this case there is one exact solution
    mwplz = w_mass*w_mass / fabs(lep.Pz());
    S1 = ( nu.Px()*nu.Px() + nu.Py()*nu.Py() - (mwplz/2.)*(mwplz/2.) ) / mwplz ;
    S2 = S1;
    Nsolution = 1;
  }
  else {
    //  use the full calculation
    B = nu.Px()*lep.Px() + nu.Py()*lep.Py();
    C1 = 1. + nu.Pt()*nu.Pt() * ( lep.Pz()*lep.Pz() - lep.E()*lep.E() ) / ((A + B)*(A + B)\
									   );
    if(C1 <= 0.) {
      // in this case C will be zero and there is only one solution,
      // so simply return the one possible solution
      S1 = (-(A + B) * lep.Pz()) / S_denom;
      S2 = S1;
      Nsolution = 1;
    }
    else {
      // finally, there are actually two solutions,
      // we have to compute them both and then choose the one closer to zero.
      C  = TMath::Sqrt(C1);
      S1 = (-( A + B) * lep.Pz() + (A + B) * lep.E()*C) / S_denom;
      S2 = (-( A + B) * lep.Pz() - (A + B) * lep.E()*C) / S_denom;
      Nsolution = 2;
    }
  }
  //by default, S1 should be the smaller one
  if ( fabs(S1) > fabs(S2) ) { //exchange S1 and S2
    double temp = S2;
    S2 = S1;
    S1 = temp;
  }
    
  return Nsolution;
}

double *ReconstructTop(const TMBLorentzVector &lep,
                       const TMBLorentzVector &jet,
                       double MET_x, double MET_y) {
  
  TMBLorentzVector nu_reco1;
  TMBLorentzVector nu_reco2;
  TMBLorentzVector Top1;
  TMBLorentzVector Top2;
  
  double S1=-999, S2=-999;
  CalculateNuPz(W_Mass, lep, MET_x, MET_y, S1, S2);
  nu_reco1.SetXYZM(MET_x, MET_y, S1, 0.0);
  nu_reco2.SetXYZM(MET_x, MET_y, S2, 0.0);
  
  Top1 = (TMBLorentzVector)MakeTop(lep, nu_reco1, jet);
  Top2 = (TMBLorentzVector)MakeTop(lep, nu_reco2, jet);
  
  if (debug)  
    printf("S1: %f S2: %f TopMass1: %f TopMass2: %f\n", 
	   (float)S1, (float)S2,
	   (float)Top1.M(), (float)Top2.M());
  
  optvar[0] = S1;
  optvar[1] = S2;
  optvar[2] = Top1.M();
  optvar[3] = Top2.M();
  optvar[4] = 0;
  optvar[5] = 0;
  
  return optvar;
}

double *ReconstructTop(const TMBLorentzVector &lep,
		       std::vector<TMBLorentzVector> &jets,
		       double MET_x, double MET_y, int ijet) {

  TMBLorentzVector nu_reco1;
  TMBLorentzVector nu_reco2;
  TMBLorentzVector Top1;
  TMBLorentzVector Top2;
  
  //double optvar[6]; //nuPz1, nuPz2, TopMass1, TopMass2 define globally already
  TH2F * h2f1 = new TH2F("h2f1", "Top Mass (solution 1) vs. neutrino p_{z}",
			 2000, -1000, 1000, 2000, -1000, 1000);
  TH2F * h2f2 = new TH2F("h2f2", "Top Mass (solution 2) vs. neutrino p_{z}",
			 2000, -1000, 1000, 2000, -1000, 1000);

  int Nentry = 1000;
  for (int i=0; i<Nentry; i++) {
    //cout << i << endl;
    double MEx_new = MET_x;
    double MEy_new = MET_y;
    double ratio = 0;
    bool   skipthis = false;
    //FixTF_Parameters( jets[0].Eta(), jets[0].E() );
    for (int jjet=0; jjet<jets.size(); jjet++) {
      double e_rnd = rndm->Gaus(jets[jjet].E(),
				fun_sigma->Eval(jets[jjet].E())*jets[jjet].E());
      if ( e_rnd < 0 && ijet!=jjet)
	e_rnd = 0;
      if ( e_rnd < 0 && ijet==jjet)
	skipthis = true;
      double tempratio = e_rnd/jets[jjet].E();
      //cout << "e_rnd: "<< e_rnd << " tempratio: " << tempratio << endl;
      if (ijet==jjet) ratio = tempratio;
      MEx_new += jets[jjet].Px() * (1 - tempratio);
      MEy_new += jets[jjet].Py() * (1 - tempratio);
    }
    if ( skipthis )
      continue;
    TMBLorentzVector jet_new(jets[ijet]);
    jet_new.SetPxPyPzE( ratio*jets[ijet].Px(),
			ratio*jets[ijet].Py(),
			ratio*jets[ijet].Pz(),
			ratio*jets[ijet].E());

    double S1=-999, S2=-999;
    CalculateNuPz(W_Mass, lep, MEx_new, MEy_new, S1, S2);
    //if (debug)  printf("S1: %f S2: %f \n", (float)S1, (float)S2);
    nu_reco1.SetXYZM(MEx_new, MEy_new, S1, 0.0);
    nu_reco2.SetXYZM(MEx_new, MEy_new, S2, 0.0);

    Top1 = (TMBLorentzVector)MakeTop(lep, nu_reco1, jet_new);
    Top2 = (TMBLorentzVector)MakeTop(lep, nu_reco2, jet_new);

    double top_mass1 = Top1.M();
    double top_mass2 = Top2.M();
    //     if (debug) {
    //       printf("==TestingNewVariables== NuPz1: %8.3f TopMass1: %8.3f\n",
    // 	     (float)nu_reco1.Pz(), (float)top_mass1);
    //       printf("==TestingNewVariables== NuPz2: %8.3f TopMass2: %8.3f\n",
    // 	     (float)nu_reco2.Pz(), (float)top_mass2);
    //     }
    h2f1->Fill((float)nu_reco1.Pz(), (float)top_mass1);
    h2f2->Fill((float)nu_reco2.Pz(), (float)top_mass2);
  }

  //printf("\n");
  char sbuf[256];
  sprintf(sbuf, "h2f1_Jet%i_energy_%.3f", ijet, (float)jets[ijet].E());
  h2f1->SetName(sbuf);
  sprintf(sbuf, "Top Mass (solution 1) vs. #nu p_{z} (Jet E_{original}: %.3f GeV)",
	  (float)jets[ijet].E());
  h2f1->SetTitle(sbuf);

  sprintf(sbuf, "h2f2_Jet%i_energy_%.3f", ijet, (float)jets[ijet].E());
  h2f2->SetName(sbuf);
  sprintf(sbuf, "Top Mass (solution 2) vs. #nu p_{z} (Jet E_{original}: %.3f GeV)",
	  (float)jets[ijet].E());
  h2f2->SetTitle(sbuf);

  //     monitor_file->cd();
  //h2f1->Write();
  //h2f2->Write();
  //     monitor_file->Close();
  //     monitor_file->Close();

  int max_bin_pz_S1 = -1;
  int max_bin_pz_S2 = -1;
  int max_bin_tm_S1 = -1;
  int max_bin_tm_S2 = -1; //tm: top mass
  int max_bin_z_tmp = -1; //should be null

  h2f1->GetMaximumBin(max_bin_pz_S1, max_bin_tm_S1, max_bin_z_tmp);
  h2f2->GetMaximumBin(max_bin_pz_S2, max_bin_tm_S2, max_bin_z_tmp);

  //6 variables are stored in an array
  optvar[0] = h2f1->ProjectionX()->GetXaxis()->GetBinCenter(max_bin_pz_S1);
  optvar[1] = h2f2->ProjectionX()->GetXaxis()->GetBinCenter(max_bin_pz_S2);
  optvar[2] = h2f1->ProjectionY()->GetXaxis()->GetBinCenter(max_bin_tm_S1);
  optvar[3] = h2f2->ProjectionY()->GetXaxis()->GetBinCenter(max_bin_tm_S2);
  optvar[4] = h2f1->GetBinContent(max_bin_pz_S1, //S1, posterior value
				  max_bin_tm_S1)/Nentry;
  optvar[5] = h2f2->GetBinContent(max_bin_pz_S2, //S2, posterior value
				  max_bin_tm_S2)/Nentry;

  delete h2f1;
  delete h2f2;

  return optvar;
}

TMBLorentzVector MakeTop(const TMBLorentzVector &lep, 
			 const TMBLorentzVector &nu, 
			 const TMBLorentzVector &jet) {
  TMBLorentzVector W;
  W.SetXYZM(0., 0., 0., 0.);
  W = lep+nu;

  TMBLorentzVector Top;
  Top.SetXYZM(0., 0., 0., 0.);
  Top = W+jet;
  return Top;
}