#include <cstdlib>
#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <string>
#include <vector>

#include "TMath.h"
#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TLatex.h"
#include "TString.h"
#include "TObjString.h"
#include <TMinuit.h>
#include <TColor.h>
#include <limits.h>
#include <float.h>
#include "TSystem.h"
#include "TROOT.h"

#include "TMVA/Factory.h"
#include "TMVA/DataLoader.h"
#include "TMVA/Reader.h"
#include "TMVA/Tools.h"
#include "TMVA/MethodCuts.h"
#include "TMVA/TMVAGui.h"

using namespace std;
using namespace TMVA;

/*--------------------------------------------------------------------------*
 *                     Function to Get TTrees from File                     *
 *--------------------------------------------------------------------------*/

TTree* getTree(TString filename, TString TupleName) {
// Get the TTree objects from files
  TFile* inputfile = new TFile(filename);
  TTree* Tree = (TTree*)inputfile->Get(TupleName);

  return Tree;

  delete inputfile;

}



/*--------------------------------------------------------------------------*
 *                          Significance Equation                           *
 *--------------------------------------------------------------------------*/

double ZA_fcn(double s, double b) {

  double Z = sqrt( 2. * ( (s+b)*TMath::Log(1 + (s/b)) - s ) );
  if ( Z > 100. ) { Z = 0.; }

  return Z;
}



/*--------------------------------------------------------------------------*
 *                 Significance Function for Minuit Analysis                *
 *--------------------------------------------------------------------------*/

// fcn passes back f = chi, the function to be minimized.
void fcn(int& nvars, double* deriv, double& f, double par[], int flag) {
  TString sigFilename = "hhttbb.root";
  TString bckFilename[] = {"ttbar_new.root", "Ztautau.root", "ZtautauB.root", "ZtautauC.root", "ZtautauL.root"};
  TString TupleName = "CxAODTuple_Nominal";
  const int nBackgrounds = 5;

// Get the TTree objects from files
  TFile* siginputfile = new TFile(sigFilename);
  TTree* sigTree = (TTree*)siginputfile->Get(TupleName);

// Get the TTree objects from files
  TFile* bckinputfile[nBackgrounds];
  TTree* bckTree[nBackgrounds];
  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    bckinputfile[i] = new TFile(bckFilename[i]);
    bckTree[i] = (TTree*)bckinputfile[i]->Get(TupleName);
  }

  TString varsName[] = {"mBB", "mMMC","mHH", "drBB", "drLepTau", "mtLepMet"};
  double sf[] = {1e-3, 1., 1e-3, 1., 1., 1e-3};

  const int nSig = sigTree->GetEntries();
  int nBck[nBackgrounds];
  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    nBck[i] = bckTree[i]->GetEntries();
  }

// set address of variable to cut over
  TString SVN = "SelVars_Nominal_";
  float var[nvars], weight;
  sigTree->SetBranchAddress("EventWeight_Nominal_TotalWeight", &weight);
  for ( int i=0 ; i<nvars ; i++ ) {
    sigTree->SetBranchAddress(SVN + varsName[i], &var[i]);
  }

  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    bckTree[i]->SetBranchAddress("EventWeight_Nominal_TotalWeight", &weight);
    for ( int j=0 ; j<nvars ; j++ ) {
      bckTree[i]->SetBranchAddress(SVN + varsName[j], &var[j]);
    }
  }

  double cutVals[nvars][2];
  for ( int i=0 ; i<nvars ; i++ ) {
    cutVals[i][0] = par[i];
    cutVals[i][1] = par[i+nvars];
  }

// get the signal events for each cut
  double sigCount = 0.;
  double sigTotalWeight = 0.;

  for ( int i=0 ; i<nSig ; i++ ) {
    sigTree->GetEvent(i);
    bool pass = true;
    for ( int k=0 ; k<nvars ; k++ ) {
      if ( var[k] < cutVals[k][0]/sf[k] ) { pass = false;}
      if ( var[k] > cutVals[k][1]/sf[k] ) { pass = false;}
    }

    if ( pass == true ) { sigCount += weight; }

    sigTotalWeight += weight;
    }

// get the background events for each cut
  double bckCount[nBackgrounds];
  for( int i=0 ; i<nBackgrounds ; i++) { bckCount[i] = 0.;}

  double bckTotalWeight[nBackgrounds];
  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    bckTotalWeight[i] = 0.;
    for ( int j=0 ; j<nBck[i] ; j++ ) {
      bckTree[i]->GetEvent(j);
      bool pass = true;
      for ( int l=0 ; l<nvars ; l++ ) {
        if ( var[l] < cutVals[l][0]/sf[l] ) { pass = false; }
        if ( var[l] > cutVals[l][1]/sf[l] ) { pass = false; }
      }

      if ( pass == true ) { bckCount[i] += weight; }

      bckTotalWeight[i] += weight;
    }
  }

// find the significance value
  double s = sigCount/sigTotalWeight;
  double b = 0.;
  for ( int j=0 ; j<nBackgrounds ; j++ ) {
    b += bckCount[j]/bckTotalWeight[j];
  }

  double Z;
  if ( b!=0. ) { Z = sqrt( 2. * ( (s+b)*TMath::Log(1 + (s/b)) - s ) );;}
  else { Z=0. ;}



  delete sigTree;
  delete siginputfile;
  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    delete bckTree[i];
    delete bckinputfile[i];
  }

  f = -1.*Z;

}




/*--------------------------------------------------------------------------*
 *                         Scan for Minuit Analysis                         *
 *--------------------------------------------------------------------------*/

void scan(double *outparB, double *outparA, int varsIndex, int nvars, TString varsName, TString varsLabel, double *xLimits, double sf) {

  double outpar[nvars][2];
  for (int i = 0; i<nvars; i++) {
    outpar[i][0] = outparB[i];
    outpar[i][1] = outparA[i];
  }

  int icondn;

// scan over variable being varied
  gMinuit->mncomd(Form("scan %i 1000 %lg %lg", varsIndex+1, xLimits[0], xLimits[1]), icondn);

  TCanvas *cScan = new TCanvas();
  TGraph *gScan = (TGraph*)gMinuit->GetPlot();
  gScan->SetTitle("Scan " + varsLabel + "; " + varsLabel + "; Significance, Z");
  //gScan->GetYaxis()->SetRangeUser(30, 40);
  gScan->Draw("al");

// Print the scan as a png
  TImage *img;
  img = TImage::Create();
  img->FromPad(cScan);
  if ( varsIndex<nvars ) {
    img->WriteImage(Form("./Cuts/Multiple Variable Cuts/%i. ", varsIndex+1) + varsName + "_Scan_t>tcut.png");
  }
  else {
    img->WriteImage(Form("./Cuts/Multiple Variable Cuts/%i. ", varsIndex+1-nvars) + varsName + "_Scan_t<tcut.png");
  }

  cout << "Minuit Scan of variable " << varsName << " complete" << endl << endl;

  delete img;
  delete cScan;

}



/*--------------------------------------------------------------------------*
 *                              Minuit Analysis                             *
 *--------------------------------------------------------------------------*/

void minimisation(const int nBackgrounds, TString *varsName, const int nvars, TString *varsLabel, TString *units, double *sf, int cutVarIndex, double *xLimits) {
  // Initialize minuit, set initial values etc. of parameters.
  TMinuit minuit(nvars*2);			// the number of parameters = nvars * 2 cuts on each
  minuit.SetFCN(fcn);

  double cutVals[nvars][2];				// the start values
  int icondn;

/*--------------------------------------------------------------
  Cut Below			  Cut above
  --------------------------------------------------------------*/
// starting value
  cutVals[0][0] = 72.;		  cutVals[0][1] = 141.;		// mBB
  cutVals[1][0] = 98.;		  cutVals[1][1] = 476.;		// mMMC
  /*cutVals[2][0] = 338.;		  cutVals[2][1] = 3.93921e+06;	// mHH
  cutVals[3][0] = 0.76;		  cutVals[3][1] = 4.7;		// drBB
  cutVals[4][0] = 0.6;		  cutVals[4][1] = 4.9;		// drLepTau
  cutVals[5][0] = 15.;	          cutVals[5][1] = 190.;		// mtLepMet*/


//minuit.DefineParameter(i, Variable Name,           Starting Value, stepSize, minVal, maxVal
  /*minuit.DefineParameter(0,  varsName[0]+"_t>tcut",  cutVals[0][0],  0.1,      30.,    100.); // mBB
  minuit.DefineParameter(1,  varsName[1]+"_t>tcut",  cutVals[1][0],  0.5,      70.,    110.); // mMMC
  minuit.DefineParameter(2,  varsName[2]+"_t>tcut",  cutVals[2][0],  0.5,      250.,   700.); // mHH
  minuit.DefineParameter(3,  varsName[3]+"_t>tcut",  cutVals[3][0],  0.001,    0.4,    2.);  // drBB
  minuit.DefineParameter(4,  varsName[4]+"_t>tcut",  cutVals[4][0],  0.001,    0.4,    2.);
  minuit.DefineParameter(5,  varsName[5]+"_t>tcut",  cutVals[5][0],  0.5,      0.,     200.);

  minuit.DefineParameter(6,  varsName[0]+"_t<tcut",  cutVals[0][1],  0.5,      130.,   500.); // mBB
  minuit.DefineParameter(7,  varsName[1]+"_t<tcut",  cutVals[1][1],  0.5,      110.,   500.); // mMMC
  minuit.DefineParameter(8,  varsName[2]+"_t<tcut",  cutVals[2][1],  0.5,      0.,     0.);   // mHH
  minuit.DefineParameter(9,  varsName[3]+"_t<tcut",  cutVals[3][1],  0.001,    1.,     5.);   // drBB
  minuit.DefineParameter(10, varsName[4]+"_t<tcut",  cutVals[4][1],  0.001,    0.4,    5.);
  minuit.DefineParameter(11, varsName[5]+"_t<tcut",  cutVals[5][1],  0.5,      0.,     200.);*/

  minuit.DefineParameter(0,  varsName[0]+"_t>tcut",  cutVals[0][0],  0.1,      30.,    100.); // mBB
  minuit.DefineParameter(1,  varsName[1]+"_t>tcut",  cutVals[1][0],  0.1,      70.,    110.); // mMMC
  minuit.DefineParameter(2,  varsName[0]+"_t<tcut",  cutVals[0][1],  0.1,      130.,   500.); // mBB
  minuit.DefineParameter(3,  varsName[1]+"_t<tcut",  cutVals[1][1],  0.1,      110.,   500.); // mMMC

  for (int i=0 ; i<(2*nvars) ; i++ ) {
   if ( i!=cutVarIndex) {gMinuit->mncomd(Form("FIX %i", i+1), icondn);} // set i!= the variable you
  }

  // Do the minimization!
  gMinuit->mncomd("SIMplex 1000 0.1", icondn);
  gMinuit->mncomd("MIGrad 20000 0.1", icondn);

  /*
  FMIN: the best function value found so far
  FEDM: the estimated vertical distance remaining to minimum
  ERRDEF: the value of UP defining parameter uncertainties
  NPARI: the number of currently variable parameters
  PARX: the highest (external) parameter number defined by user
  ISTAT: a status integer indicating how good is the covariance matrix:
  0= not calculated at all
  1= approximation only, not accurate
  2= full matrix, but forced positive-definite
  3= full accurate covariance matrix
  */

  double fmin, fedm, errdef;
  int npari, nparx, istat;

  minuit.mnstat(fmin, fedm, errdef, npari, nparx, istat);

  cout << "fmin = " << fmin << " fedm = " << fedm << " errdef = " << errdef << endl;
  cout << "npari = " << npari << " nparx = " << nparx << " istat = " << istat << endl;

  double outpar[nvars][2], err[nvars][2];
  for (int i = 0; i<nvars; i++) {
    gMinuit->GetParameter(i, outpar[i][0], err[i][0]);
    gMinuit->GetParameter(i+nvars, outpar[i][1], err[i][1]);
  }

  cout << "minimisation of variable " << varsName[cutVarIndex] << " complete" << endl << endl;

  double minResults[25]; //4*nvars +1

  minResults[0] = fmin;
  for (int i = 0; i<nvars; i++) {
    minResults[4*i +1] = outpar[i][0];
    minResults[4*i +2] = err[i][0];
    minResults[4*i +3] = outpar[i][1];
    minResults[4*i +4] = err[i][1];
  }

  double outparB[nvars], outparA[nvars];
  for (int i = 0; i<nvars; i++) {
    outparB[i] = outpar[i][0];
    outparA[i] = outpar[i][1];
  }

  if (cutVarIndex<nvars) {
    scan(outparB, outparA, cutVarIndex, nvars, varsName[cutVarIndex], varsLabel[cutVarIndex], xLimits, sf[cutVarIndex]);
  } else {
    scan(outparB, outparA, cutVarIndex, nvars, varsName[cutVarIndex-nvars], varsLabel[cutVarIndex-nvars], xLimits, sf[cutVarIndex-nvars]);
  }

}



void Debugger2() {

/*--------------------------------------------------------------------------*
 *                    Set up basic data for all variables                   *
 *--------------------------------------------------------------------------*/

  TString sigFilename = "hhttbb.root";
  TString bckFilename[] = {"ttbar_new.root", "Ztautau.root", "ZtautauB.root", "ZtautauC.root", "ZtautauL.root"};
  TString NTupleName = "CxAODTuple_Nominal";
  const int nBackgrounds = 5;

  TTree* sigTree = getTree(sigFilename, NTupleName);
  TTree *bckTree[nBackgrounds];
  for ( int i=0 ; i<nBackgrounds ; i++ ) {
    bckTree[i] = getTree(bckFilename[i], NTupleName);
  }

  TString varsName[] = {"mBB", "mMMC"/*,"mHH", "drBB", "drLepTau", "mtLepMet"*/};
  const int nvars = 2;
  
  TString varsLabel[] = {"m_{bb}", "m^{MMC}_{#tau#tau}"/*, "m_{hh}", "#DeltaR(b,b)", "#DeltaR(#tau,#tau)", "m^{W}_{T}"*/};

  TString units[] = {"[GeV]", "[GeV]"/*, " [GeV]", " ", " ", "[GeV]"*/};

  double sf[] = {1e-3, 1., 1e-3, 1., 1., 1e-3};
  //double Xlimits[nvars][2] = {{0., 500.}, {0., 500.}, {200., 1200.}, {0., 5}, {0., 5}, {0., 200.}};
  double Xlimits[nvars][2] = {{0., 500.}, {0., 500.}};



/*--------------------------------------------------------------------------*
 *               Get graphs of optimised cuts for all variables             *
 *--------------------------------------------------------------------------*/

  for ( int i=0 ; i<nvars ; i++){
    if (i!=3 && i!=4 && i!=5) {minimisation(nBackgrounds, varsName, nvars, varsLabel, units, sf, i, Xlimits[i]);}

    if (i!=2) {minimisation(nBackgrounds, varsName, nvars, varsLabel, units, sf, i+nvars, Xlimits[i]);}
  }

}