#include "TH1D.h"
#include "TH2D.h"
#include "THStack.h"
#include "TGraph.h"
#include "TTree.h"
#include "TChain.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TStyle.h"
#include "TColor.h"
#include "TLorentzVector.h"
#include "TRandom3.h"
#include "TLorentzVector.h"
#include "TString.h"

#include "RooDataHist.h"
#include "RooRealVar.h"
#include "RooPlot.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooRealConstant.h"
#include "RooAddPdf.h"
#include "RooFitResult.h"
#include "RooProduct.h"
#include "RooCrystalBall.h"
#include "RooBreitWigner.h"
#include "RooArgSet.h"
#include "RooFFTConvPdf.h"
#include "RooNovosibirsk.h"

#include <string>
#include <iostream>
#include <cmath>
#include <vector>

template <typename T>
class Cut
{
public:
  virtual bool Check(T y) = 0;
  virtual std::string ToString() = 0;
};

template <typename T>
class LT : public Cut<T>
{
private:
  T value;

public:
  LT(T val) : value{val}
  {
  }

  bool Check(T y)
  {
    return y < value;
  }

  std::string ToString()
  {
    return TString::Format("(x < %s)", std::to_string(value).c_str()).Data();
  }
};

template <typename T>
class GT : public Cut<T>
{
private:
  T value;

public:
  GT(T val) : value{val} {}

  bool Check(T y)
  {
    return y > value;
  }

  std::string ToString()
  {
    return TString::Format("(x > %s)", std::to_string(value).c_str()).Data();
  }
};

template <typename T>
class EQ : public Cut<T>
{
private:
  T value;

public:
  EQ(T val) : value{val} {}

  bool Check(T y)
  {
    return y == value;
  }

  std::string ToString()
  {
    return TString::Format("(x == %s)", std::to_string(value).c_str()).Data();
  }
};

template <typename T>
class BT : public Cut<T>
{
private:
  T low;
  T up;

public:
  BT(T lower, T upper) : low{lower}, up{upper} {}

  bool Check(T y)
  {
    return low < y && y < up;
  }

  std::string ToString()
  {
    return TString::Format("(x in [%s, %s])", std::to_string(low).c_str(), std::to_string(up).c_str()).Data();
  }
};

template <typename T>
class Variable
{
private:
  const char *name;
  std::vector<Cut<T> *> cuts;

  T value;

public:
  Variable()
  {
  }

  Variable(const char *name, std::vector<Cut<T> *> cuts)
  {
    this->name = name;
    this->cuts = cuts;
  }

  T *GetValuePtr()
  {
    return &value;
  }

  T GetValue() const
  {
    return value;
  }

  bool CheckCuts() const
  {
    for (int i = 0; i < cuts.size(); i++)
    {
      if (!cuts[i]->Check(value))
      {
        return false;
      }
    }
    return true;
  }
};

class VariableCollection
{
private:
  std::map<std::string, Variable<Double_t>> doubles;
  std::map<std::string, Variable<Float_t>> floats;
  std::map<std::string, Variable<Int_t>> ints;
  std::map<std::string, Variable<Bool_t>> bools;
  std::vector<std::vector<std::string>> trigger_cuts;

  TLorentzVector ComputeFourMomentum(const char *name, bool do_mass_sub, double subbed_mass)
  {
    std::string pxname = TString::Format("%s_PX", name).Data();
    std::string pyname = TString::Format("%s_PY", name).Data();
    std::string pzname = TString::Format("%s_PZ", name).Data();
    std::string ename = TString::Format("%s_ENERGY", name).Data();

    auto itx = floats.find(pxname);
    if (itx == floats.end())
    {
      return TLorentzVector{};
    }

    auto ity = floats.find(pyname);
    if (ity == floats.end())
    {
      return TLorentzVector{};
    }

    auto itz = floats.find(pzname);
    if (itz == floats.end())
    {
      return TLorentzVector{};
    }

    TVector3 momentum(itx->second.GetValue(), ity->second.GetValue(), itz->second.GetValue());
    double energy = 0;

    if (do_mass_sub)
    {
      energy = TMath::Sqrt(TMath::Sq(subbed_mass) + momentum.Mag2());
    }
    else
    {
      auto ite = floats.find(ename);
      if (ite == floats.end())
      {
        return TLorentzVector{};
      }
      energy = ite->second.GetValue();
    }

    return TLorentzVector(momentum, energy);
  }

public:
  VariableCollection()
  {
    doubles = std::map<std::string, Variable<Double_t>>{};
    floats = std::map<std::string, Variable<Float_t>>{};
    ints = std::map<std::string, Variable<Int_t>>{};
    bools = std::map<std::string, Variable<Bool_t>>{};
    trigger_cuts = std::vector<std::vector<std::string>>{};
  }

  void Connect(TChain *chain)
  {
    for (auto &[key, var] : doubles)
    {
      chain->SetBranchAddress(key.c_str(), var.GetValuePtr());
    }

    for (auto &[key, var] : floats)
    {
      chain->SetBranchAddress(key.c_str(), var.GetValuePtr());
    }

    for (auto &[key, var] : ints)
    {
      chain->SetBranchAddress(key.c_str(), var.GetValuePtr());
    }

    for (auto &[key, var] : bools)
    {
      chain->SetBranchAddress(key.c_str(), var.GetValuePtr());
    }
  }

  void AddDouble(const char *name, const std::vector<Cut<Double_t> *> &cuts)
  {
    doubles[name] = Variable<Double_t>(name, cuts);
  }

  void AddFloat(const char *name, const std::vector<Cut<Float_t> *> &cuts)
  {
    floats[name] = Variable<Float_t>(name, cuts);
  }

  void AddInt(const char *name, const std::vector<Cut<Int_t> *> &cuts)
  {
    ints[name] = Variable<Int_t>(name, cuts);
  }

  void AddBool(const char *name, const std::vector<Cut<Bool_t> *> &cuts)
  {
    bools[name] = Variable<Bool_t>(name, cuts);
  }

  void AddTriggerCut(const std::vector<std::string> &triggers_list)
  {
    for (const auto &trig : triggers_list)
    {
        bools[trig.c_str()] = Variable<Bool_t>(trig.c_str(), {});
    }
    trigger_cuts.push_back(triggers_list);
  }

  void AddFourMomentumFor(const char *name)
  {
    std::string pxname = TString::Format("%s_PX", name).Data();
    std::string pyname = TString::Format("%s_PY", name).Data();
    std::string pzname = TString::Format("%s_PZ", name).Data();
    std::string ename = TString::Format("%s_ENERGY", name).Data();

    floats[pxname] = Variable<Float_t>(pxname.c_str(), {});
    floats[pyname] = Variable<Float_t>(pyname.c_str(), {});
    floats[pzname] = Variable<Float_t>(pzname.c_str(), {});
    floats[ename] = Variable<Float_t>(ename.c_str(), {});
  }

  Double_t GetDouble(const char *name) const
  {
    auto it = doubles.find(name);
    if (it != doubles.end())
    {
      return it->second.GetValue();
    }
    return 0;
  }

  Float_t GetFloat(const char *name) const
  {
    auto it = floats.find(name);
    if (it != floats.end())
    {
      return it->second.GetValue();
    }
    return 0;
  }

  Int_t GetInt(const char *name) const
  {
    auto it = ints.find(name);
    if (it != ints.end())
    {
      return it->second.GetValue();
    }
    return 0;
  }

  Bool_t GetBool(const char *name) const
  {
    auto it = bools.find(name);
    if (it != bools.end())
    {
      return it->second.GetValue();
    }
    return 0;
  }

  TLorentzVector GetFourMomentum(const char *name)
  {
    return ComputeFourMomentum(name, false, 0.);
  }

  TLorentzVector GetFourMomentumWithMassSub(const char *name, double subbed_mass)
  {
    return ComputeFourMomentum(name, true, subbed_mass);
  }

  bool CheckCuts() const
  {
    for (const auto &[key, var] : doubles)
    {
      if (!var.CheckCuts())
      {
        return false;
      }
    }

    for (const auto &[key, var] : floats)
    {
      if (!var.CheckCuts())
      {
        return false;
      }
    }

    for (const auto &[key, var] : ints)
    {
      if (!var.CheckCuts())
      {
        return false;
      }
    }

    for (const auto &[key, var] : bools)
    {
      if (!var.CheckCuts())
      {
        return false;
      }
    }

    for (const auto &triggers : trigger_cuts)
    {
      bool cut_okay = false;
      for (const auto &trigger : triggers)
      {
        cut_okay = cut_okay | this->GetBool(trigger.c_str());
      }

      if (!cut_okay)
      {
        return false;
      }
    }

    return true;
  }

  void PrintValues() const
  {
    for (const auto &[key, var] : doubles)
    {
      std::cout << key << ": " << var.GetValue() << std::endl;
    }

    for (const auto &[key, var] : floats)
    {
      std::cout << key << ": " << var.GetValue() << std::endl;
    }

    for (const auto &[key, var] : ints)
    {
      std::cout << key << ": " << var.GetValue() << std::endl;
    }

    for (const auto &[key, var] : bools)
    {
      std::cout << key << ": " << var.GetValue() << std::endl;
    }
  }
};

void BuToHpMuMu()
{

  TChain *data_chain = new TChain("BuToHpMuMu/DecayTree");
  data_chain->Add("/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0r1_tuple_90000000_2023_v0r0p6288631.root");
  data_chain->Add("/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0_tuple_90000000_v0r0p6288631.root");

  VariableCollection col{};
  col.AddFourMomentumFor("L1");
  col.AddFourMomentumFor("L2");
  col.AddFourMomentumFor("Hp");

  col.AddTriggerCut({"Hlt2_InclDetDiMuonDecision", "Hlt2_InclDetDiMuon_3BodyDecision", "Hlt2_InclDetDiMuon_4BodyDecision"});
  col.AddTriggerCut({"Hlt1TrackMVADecision", "Hlt1TwoTrackMVADecision"});

  col.Connect(data_chain);

  const Double_t MASS_HIST_MIN = 5100.;
  const Double_t MASS_HIST_MAX = 6000.;
  const int N_BINS = 70;
  TH1D *h1_B_M = new TH1D("h1_B_M", "B Mass", N_BINS, MASS_HIST_MIN, MASS_HIST_MAX);

  unsigned int entries = data_chain->GetEntries();
  for (unsigned int i = 0; i < entries; i++)
  {
    data_chain->GetEntry(i);

    if (col.CheckCuts())
    {
      auto l1_4v = col.GetFourMomentum("L1");
      auto l2_4v = col.GetFourMomentum("L2");
      auto K_4v = col.GetFourMomentumWithMassSub("Hp", 493.677);

      h1_B_M->Fill((l1_4v + l2_4v + K_4v).M());
    }

    if ((i + 1) % 10000 == 0 || i + 1 == entries)
    {
      std::cout << "["
                << "BuToHpMuMu"
                << "] Processed event: " << i + 1 << " (" << TString::Format("%.2f", ((double)(i + 1) / (double)entries) * 100.) << "%)" << std::endl;
    }
  }

  TCanvas *c1 = new TCanvas("c1", "c1", 0, 0, 800, 600);
  h1_B_M->Draw();
  c1->Draw();

  return 0;
}

int do_all() {
  BuToHpMuMu();
  return 0;
}