#ifndef HLT1_DECISION_ANALYSIS
#define HLT1_DECISION_ANALYSIS

#include <string>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <filesystem>
#include <string_view>


struct Hlt1Decision
{
  std::string name;
  int index;
  Bool_t value;

  std::string GetName() const
  {
    return TString::Format("Hlt1%sDecision", name.c_str()).Data();
  }

  Bool_t *GetValuePointer()
  {
    return &value;
  }
};

std::vector<Hlt1Decision> Hlt1Decisions{
    Hlt1Decision{"DiMuonHighMass", 1},
    Hlt1Decision{"DiMuonLowMass", 2},
    Hlt1Decision{"DiMuonNoIP", 3},
    Hlt1Decision{"DiMuonSoft", 4},
    Hlt1Decision{"DisplacedLeptons", 5},
    Hlt1Decision{"LowPtDiMuon", 6},
    Hlt1Decision{"LowPtMuon", 7},
    Hlt1Decision{"OneMuonTrackLine", 8},
    Hlt1Decision{"SingleHighEt", 9},
    Hlt1Decision{"SingleHighPtMuon", 10},
    Hlt1Decision{"TrackMVA", 11},
    Hlt1Decision{"TrackMuonMVA", 12},
    Hlt1Decision{"TwoTrackMVA", 13},
};

struct Hlt1DecisionPlot
{
  std::string name;
  bool exclusive;
  std::set<std::string> lines;
};

const std::vector<Hlt1DecisionPlot>
    Hlt1DecisionSets{
        Hlt1DecisionPlot{"mva_excl", true, std::set<std::string>{"TwoTrackMVA", "TrackMuonMVA", "TrackMVA"}},

        Hlt1DecisionPlot{"mva_incl", false, std::set<std::string>{"TwoTrackMVA", "TrackMuonMVA", "TrackMVA"}},

        Hlt1DecisionPlot{"pt_excl", true, std::set<std::string>{"LowPtMuon", "LowPtDiMuon", "DisplacedLeptons"}},

        Hlt1DecisionPlot{"pt_incl", false, std::set<std::string>{"LowPtMuon", "LowPtDiMuon", "DisplacedLeptons"}},

        Hlt1DecisionPlot{"dimu_excl", true, std::set<std::string>{"DiMuonNoIP", "DiMuonLowMass", "DiMuonHighMass"}},

        Hlt1DecisionPlot{"dimu_incl", false, std::set<std::string>{"DiMuonNoIP", "DiMuonLowMass", "DiMuonHighMass"}},

        Hlt1DecisionPlot{"mvapt_incl", false, std::set<std::string>{"TwoTrackMVA", "TrackMuonMVA", "TrackMVA", "LowPtMuon", "LowPtDiMuon", "DisplacedLeptons"}},

        Hlt1DecisionPlot{"mvadimu_incl", false, std::set<std::string>{"TwoTrackMVA", "TrackMuonMVA", "TrackMVA", "DiMuonNoIP", "DiMuonLowMass", "DiMuonHighMass"}},

        Hlt1DecisionPlot{"mvadimupt_incl", false, std::set<std::string>{"TwoTrackMVA", "TrackMuonMVA", "TrackMVA", "LowPtMuon", "LowPtDiMuon", "DisplacedLeptons", "DiMuonNoIP", "DiMuonLowMass", "DiMuonHighMass"}},

        Hlt1DecisionPlot{"ptdimu_incl", false, std::set<std::string>{"LowPtMuon", "LowPtDiMuon", "DisplacedLeptons", "DiMuonNoIP", "DiMuonLowMass", "DiMuonHighMass"}},

    };

bool CutHlt1DecisionsAnd(const std::set<std::string> &decision_list)
{
  bool okay = true;
  for (const auto &var : Hlt1Decisions)
  {
    if (decision_list.find(var.name) != decision_list.end())
    {
      okay = okay && var.value;
    }
  }
  return okay;
}

bool CutHlt1DecisionsOr(const std::set<std::string> &decision_list)
{
  bool okay = false;
  for (const auto &var : Hlt1Decisions)
  {
    if (decision_list.find(var.name) != decision_list.end())
    {
      okay = okay || var.value;
    }
  }
  return okay;
}

bool CutHlt1DecisionsOrOnly(const std::set<std::string> &decision_list)
{
  bool okay = false;
  for (const auto &var : Hlt1Decisions)
  {
    if (decision_list.find(var.name) != decision_list.end())
    {
      okay = okay || var.value;
    }
    else
    {
      if (var.value)
      {
        okay = false;
        break;
      }
    }
  }
  return okay;
}

void ConnectHlt1Decisions(TChain *chain, TH2D *incl_hist, TH2D *excl_hist)
{
  for (auto &var : Hlt1Decisions)
  {
    if (chain->FindBranch(var.GetName().c_str())) {
      chain->SetBranchAddress(var.GetName().c_str(), var.GetValuePointer());
      incl_hist->Fill(0., var.name.c_str(), 0.);
      excl_hist->Fill(0., var.name.c_str(), 0.);
    }
  }
}

void CheckHlt1Decisioins(TH2D *incl_hist, TH2D *excl_hist, std::map<std::string, int> &exclusive_hits, const double reco_mass)
{
  for (const auto &var : Hlt1Decisions)
  {
    if (var.value)
    {
      incl_hist->Fill(reco_mass, var.name.c_str(), 1);
    }
  }

  bool exclusive = true;
  std::string line{};
  for (const auto &var : Hlt1Decisions)
  {
    if (var.value)
    {
      if (!line.empty())
      {
        exclusive = false;
      }

      line = var.name;
    }
  }

  if (!line.empty() && exclusive)
  {
    int &hits = exclusive_hits[line];
    if (hits)
    {
      hits += 1;
    }
    else
    {
      hits = 1;
    }
    excl_hist->Fill(reco_mass, line.c_str(), 1);
  }
}

std::vector<TH1D *> CreateHlt1DecisionHistos(const char *analysis_name)
{
  std::vector<TH1D *> histos{};
  for (int i = 0; i < Hlt1DecisionSets.size(); i++)
  {
    TH1D *h1_B_Mass_hlt1 = new TH1D(TString::Format("h1_B_Mass_hlt1_%s", Hlt1DecisionSets[i].name.c_str()), TString::Format("(%s) (%s)", Hlt1DecisionSets[i].name.c_str(), analysis_name), B_MASS_HIST_BINS, B_MASS_VAR_MIN, B_MASS_VAR_MAX);
    histos.push_back(h1_B_Mass_hlt1);
  }

  return histos;
}

void FillHlt1DecisionHistos(std::vector<TH1D *> histos, double reco_mass)
{
  for (int i = 0; i < Hlt1DecisionSets.size(); i++)
  {
    if (Hlt1DecisionSets[i].exclusive)
    {
      if (CutHlt1DecisionsOrOnly(Hlt1DecisionSets[i].lines))
      {
        histos[i]->Fill(reco_mass);
      }
    }
    else
    {
      if (CutHlt1DecisionsOr(Hlt1DecisionSets[i].lines))
      {
        histos[i]->Fill(reco_mass);
      }
    }
  }
}

void DrawHlt1DecisionHistos(const char *folder, std::vector<TH1D *> histos)
{
  std::filesystem::create_directory(TString::Format("output_files/analysis/%s", folder).Data());
  TString name = TString::Format("%s_canvas", "hlt1_decisions");
  TCanvas *c = new TCanvas(name, "HLT 1 Decisions", 0, 0, 1200, 700);
  c->Divide(4, 3);
  for (int i = 0; i < histos.size(); i++)
  {
    c->cd(i + 1);
    histos[i]->SetStats(0);
    histos[i]->Draw();

    TLegend *leg1 = new TLegend(0.58, 0.89 - (0.05 * (Hlt1DecisionSets[i].lines.size() + (int)(Hlt1DecisionSets[i].lines.size() / 3))), 0.88, 0.89);
    leg1->SetFillStyle(0);
    leg1->SetLineStyle(0);
    leg1->SetMargin(0.1);
    int index = 1;
    std::for_each(Hlt1DecisionSets[i].lines.begin(), Hlt1DecisionSets[i].lines.end(), [leg1, i, &index](std::string s)
                  { 
                    leg1->AddEntry((TObject *)0, s.c_str(), "");
                    if (index != Hlt1DecisionSets[i].lines.size() && index % 3 == 0) {
                      leg1->AddEntry((TObject *)0, "", "");
                    }
                    index++; });
    leg1->AddEntry((TObject *)0, TString::Format("# Entries: %d", (int)histos[i]->GetEntries()), "");
    leg1->Draw();
  }
  c->Draw();
  c->SaveAs(TString::Format("output_files/analysis/%s/%s.pdf", folder, name.Data()).Data());
}

#endif