#include <filesystem>

#include "TMVA/efficiencies.h"

#include "TH2F.h"
#include "TFile.h"
#include "TIterator.h"
#include "TKey.h"

void set_tmva_style()
{
  TStyle *TMVAStyle = gROOT->GetStyle("TMVA");
  if (TMVAStyle != 0)
  {
    gROOT->SetStyle("TMVA");
    return;
  }

  TMVAStyle = new TStyle(*gROOT->GetStyle("Plain")); // our style is based on Plain
  TMVAStyle->SetName("TMVA");
  TMVAStyle->SetTitle("TMVA style based on \"Plain\" with modifications defined in tmvaglob.C");
  gROOT->GetListOfStyles()->Add(TMVAStyle);
  gROOT->SetStyle("TMVA");

  TMVAStyle->SetLineStyleString(5, "[52 12]");
  TMVAStyle->SetLineStyleString(6, "[22 12]");
  TMVAStyle->SetLineStyleString(7, "[22 10 7 10]");

  // the pretty color palette of old
  TMVAStyle->SetPalette((TMVA::gConfig().fVariablePlotting.fUsePaperStyle ? 18 : 1), 0);

  TMVAStyle->SetOptFile(0);

  // use plain black on white colors
  TMVAStyle->SetFrameBorderMode(0);
  TMVAStyle->SetCanvasBorderMode(0);
  TMVAStyle->SetPadBorderMode(0);
  TMVAStyle->SetPadColor(0);
  TMVAStyle->SetFillStyle(0);

  TMVAStyle->SetLegendBorderSize(0);

  // title properties
  // TMVAStyle->SetTitleW(.4);
  // TMVAStyle->SetTitleH(.10);
  // MVAStyle->SetTitleX(.5);
  // TMVAStyle->SetTitleY(.9);
  TMVAStyle->SetTitleFillColor(TColor::GetColor("#ffffff"));
  TMVAStyle->SetTitleTextColor(TColor::GetColor("#000000"));
  TMVAStyle->SetTitleBorderSize(0);
  TMVAStyle->SetLineColor(TMVA::TMVAGlob::getTitleBorder());
  if (!TMVA::gConfig().fVariablePlotting.fUsePaperStyle)
  {
    TMVAStyle->SetFrameFillColor(TColor::GetColor("#ffffff"));
    TMVAStyle->SetCanvasColor(TColor::GetColor("#ffffff"));
  }

  // set the paper & margin sizes
  TMVAStyle->SetPaperSize(20, 26);
  TMVAStyle->SetPadTopMargin(0.08);
  TMVAStyle->SetPadRightMargin(0.01);
  TMVAStyle->SetPadBottomMargin(0.01);
  TMVAStyle->SetPadLeftMargin(0.12);

  // use bold lines and markers
  TMVAStyle->SetMarkerStyle(21);
  TMVAStyle->SetMarkerSize(0.3);
  TMVAStyle->SetHistLineWidth(2);
  TMVAStyle->SetLineStyleString(2, "[12 12]"); // postscript dashes

  // do not display any of the standard histogram decorations
  TMVAStyle->SetOptTitle(1);
  TMVAStyle->SetTitleH(0.052);

  TMVAStyle->SetOptStat(0);
  TMVAStyle->SetOptFit(0);

  // put tick marks on top and RHS of plots
  TMVAStyle->SetPadTickX(1);
  TMVAStyle->SetPadTickY(1);
}

void set_title_vis(bool vis)
{
  TStyle *TMVAStyle = gROOT->GetStyle("TMVA");
  TMVAStyle->SetTitleTextColor(TColor::GetColor(vis ? "#000000" : "#ffffff"));
}

void find_and_resize_canvas(const char *name, const char *save_folder_name, const char *save_file_name)
{
  auto canvs = (TCanvas *)gROOT->FindObject(name);
  if (canvs != 0)
  {
    std::cout << "#### Resizing Canvas: " << canvs->GetTitle() << std::endl;
    canvs->SetCanvasSize(canvs->GetWindowWidth() * 1.5, canvs->GetWindowHeight() * 1.5);
    canvs->Update();
    canvs->SaveAs(TString::Format("output_files/tmva/%s/%s.pdf", save_folder_name, save_file_name));
  }
}

void run(const char *trigger) {
  std::filesystem::create_directory(TString::Format("output_files/tmva/%s", trigger).Data());

  TString input_file = TString::Format("%s_tmva_out.root", trigger);
  TString directory = TString::Format("%s_dataloader", trigger);

  set_title_vis(true);
  TMVA::efficiencies(directory, input_file, 2, true);
  find_and_resize_canvas("c", trigger, "efficiencies");
  
  TMVA::mvas(directory, input_file, TMVA::kMVAType, true);
  find_and_resize_canvas("canvas1", trigger, "class_output_testsamp");

  TMVA::mvas(directory, input_file, TMVA::kCompareType, true);
  find_and_resize_canvas("canvas1", trigger, "class_output_overtrain");

  set_title_vis(false);
  TMVA::variables(directory, input_file, "InputVariables_Id", "Input variables Id-transformed (training sample)", false, true);
  for (int i = 0; i < 3; i++)
  {
    find_and_resize_canvas(TString::Format("canvas%d", i + 1), trigger, TString::Format("variables%d", i + 1));
  }
}


int tmva_plotting()
{
  set_tmva_style();
  run("BuToHpMuMu");
  run("B0ToHpHmMuMu");

  return 0;
}

/*
//title properties

*/