#ifndef BASIC_ANALYSIS
#define BASIC_ANALYSIS

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

#include "TTree.h"

#include "RooHist.h"

#include "constants.h"

class FourVect
{
private:
  Float_t px, py, pz, energy;

  Float_t *PtrPX()
  {
    return &px;
  }

  Float_t *PtrPY()
  {
    return &py;
  }

  Float_t *PtrPZ()
  {
    return &pz;
  }

  Float_t *PtrEnergy()
  {
    return &energy;
  }

public:
  static FourVect *Init(TTree *tree, const char *particle)
  {
    FourVect *v = new FourVect{};
    tree->SetBranchAddress(TString::Format("%s_PX", particle).Data(), v->PtrPX());
    tree->SetBranchAddress(TString::Format("%s_PY", particle).Data(), v->PtrPY());
    tree->SetBranchAddress(TString::Format("%s_PZ", particle).Data(), v->PtrPZ());
    tree->SetBranchAddress(TString::Format("%s_ENERGY", particle).Data(), v->PtrEnergy());
    return v;
  }

  TLorentzVector LorentzVector()
  {
    return TLorentzVector(px, py, pz, energy);
  }

  TLorentzVector LorentzVector(double sub_mass)
  {
    TVector3 momentum(px, py, pz);
    float energy = TMath::Sqrt(TMath::Sq(sub_mass) + momentum.Mag2());
    return TLorentzVector(momentum, energy);
  }

  std::string ToString()
  {
    return TString::Format("(%f, %f, %f, %f)", px, py, pz, energy).Data();
  }
};

struct FittedParam
{
  std::string title;
  std::string name;
  double value;
  double error;
  bool at_limit;
  bool constant;
  int decimals;

  FittedParam(RooRealVar var, int decimals)
  {
    this->title = var.GetTitle();
    this->name = var.GetName();
    double val = var.getVal();
    this->value = val;
    this->constant = var.isConstant();
    if (!this->constant)
    {
      this->error = var.getError();
      this->at_limit = ((val == var.getMin()) || (val == var.getMax()));
    }
    else
    {
      this->error = 0.;
    }
    this->decimals = decimals;
  }

  FittedParam(std::string name, std::string title, double value, double err, int decimals) {
    this->title = title;
    this->name = name;
    this->value = value;
    this->error = err;
    this->decimals = decimals;
    this->constant = false;
    this->at_limit = false;
  }

  std::string ToString() const
  {
    if (this->constant)
    {
      return TString::Format("%s = %.*f", this->title.c_str(), this->decimals, this->value).Data();
    }
    else
    {
      return TString::Format("%s = %.*f #pm %.*f", this->title.c_str(), this->decimals, this->value, this->decimals, this->error).Data();
    }
  }
};

struct ShapeParamters
{
  double alpha_left;
  double n_left;
  double alpha_right;
  double n_right;
};

struct RooFitSummary
{
  RooPlot *fit_histogram;
  RooPlot *pull_histogram;
  std::map<std::string, std::string> pdf_names;
  std::vector<FittedParam> fitted_params;
  ShapeParamters shape_parameters;
};

void DrawInDefaultCanvas(TH1 *histogram, const char *folder, double margin_left = 0.1, Option_t *option = "")
{
  std::filesystem::create_directory(TString::Format("output_files/analysis/%s", folder).Data());
  TString name = TString::Format("%s_canvas", histogram->GetName());
  TCanvas *c = new TCanvas(name, histogram->GetName(), 0, 0, 800, 600);
  c->SetLeftMargin(margin_left);
  histogram->SetStats(0);
  histogram->Draw(option);
  c->Draw();
  c->SaveAs(TString::Format("output_files/analysis/%s/%s.pdf", folder, name.Data()).Data());
}

void DrawInDefaultCanvasStacked(std::vector<TH1D *> histograms, std::vector<Color_t> colors, std::vector<Style_t> fill_style, const char *folder, double margin_left = 0.1, Option_t *option = "")
{
  std::filesystem::create_directory(TString::Format("output_files/analysis/%s", folder).Data());
  TString name = TString::Format("%s_stack_canvas", histograms[0]->GetName());
  TCanvas *c = new TCanvas(name, histograms[0]->GetName(), 0, 0, 800, 600);
  c->SetLeftMargin(margin_left);
  std::string drwopt_2 = std::string(option).empty() ? "SAME HIST" : TString::Format("%s SAME HIST", option).Data();
  for (size_t i = 0; i < histograms.size(); i++)
  {
    const Double_t scaling_factor = 1.;
    auto hist_clone = (TH1*)histograms[i]->Clone(TString::Format("%s_clone", histograms[i]->GetName()));
    hist_clone->Scale(scaling_factor / histograms[i]->GetMaximum());
    hist_clone->SetLineColor(colors[i]);
    hist_clone->SetMaximum(scaling_factor + (scaling_factor * 0.05));
    hist_clone->SetMinimum(0.);
    hist_clone->SetStats(0);
    if (fill_style[i] != 0) {
      hist_clone->SetFillStyle(fill_style[i]);
      hist_clone->SetFillColor(colors[i]);
    }

    if (i > 0)
    {
      hist_clone->Draw(drwopt_2.c_str());
    }
    else
    {
      hist_clone->Draw(drwopt_2.c_str());
    }
  }

  c->BuildLegend(0.55, 0.79, 0.87, 0.89);

  c->Draw();
  c->SaveAs(TString::Format("output_files/analysis/%s/%s.pdf", folder, name.Data()).Data());
}

void DrawInDefaultCanvas(RooPlot *rooPlot, const char *folder, double margin_left = 0, Option_t *option = "")
{
  std::filesystem::create_directory(TString::Format("output_files/analysis/%s", folder).Data());
  TString name = TString::Format("%s_canvas", rooPlot->GetName());
  TCanvas *c = new TCanvas(name, rooPlot->GetName(), 0, 0, 800, 600);
  c->SetLeftMargin(margin_left);
  rooPlot->Draw(option);
  c->Draw();
  c->SaveAs(TString::Format("output_files/analysis/%s/%s.pdf", folder, name.Data()).Data());
}

void PrintStyles(RooPlot *plt)
{
  auto Xaxis = plt->GetXaxis();
  auto Yaxis = plt->GetYaxis();
  std::cout << "## Styles for plot " << plt->GetName() << std::endl;
  std::cout << "## X axis\n"
            << "Label Size: " << Xaxis->GetLabelSize() << ",\n"
            << "Label Offset: " << Xaxis->GetLabelOffset() << ",\n"
            << "Title Size: " << Xaxis->GetTitleSize() << ",\n"
            << "Title Offset: " << Xaxis->GetTitleOffset() << "." << std::endl;

  std::cout << "## Y axis\n"
            << "Label Size: " << Yaxis->GetLabelSize() << ",\n"
            << "Label Offset: " << Yaxis->GetLabelOffset() << ",\n"
            << "Title Size: " << Yaxis->GetTitleSize() << ",\n"
            << "Title Offset: " << Yaxis->GetTitleOffset() << "." << std::endl;
}

void DrawInDefaultCanvas(RooFitSummary fitSummary, const char *folder)
{
  std::filesystem::create_directory(TString::Format("output_files/analysis/%s", folder).Data());
  TString name = TString::Format("%s_canvas", fitSummary.fit_histogram->GetName());
  TCanvas *c = new TCanvas(name, fitSummary.fit_histogram->GetTitle(), 0, 0, 800, 600);

  Double_t xlow, ylow, xup, yup;
  c->GetPad(0)->GetPadPar(xlow, ylow, xup, yup);
  c->Divide(1, 2);

  Double_t upPad_ylow = ylow + 0.25 * (yup - ylow);
  Double_t dwPad_yup = ylow + 0.25 * (yup - ylow);

  TVirtualPad *upPad = c->GetPad(1);
  upPad->SetPad(xlow, upPad_ylow, xup, yup);

  TVirtualPad *dwPad = c->GetPad(2);
  dwPad->SetPad(xlow, ylow, xup, dwPad_yup);
  dwPad->SetTopMargin(0);

  Double_t upPad_area = (yup - upPad_ylow) * (xup - xlow);
  Double_t dwPad_area = (dwPad_yup - ylow) * (xup - xlow);

  // std::cout << "### UP AREA: " << upPad_area << " LOW AREA: " << dwPad_area << std::endl;

  const Double_t pull_title_size = 0.09;
  const Double_t pull_label_size = 0.09;

  Double_t fit_title_size = (pull_title_size * dwPad_area) / upPad_area;
  Double_t fit_label_size = (pull_label_size * dwPad_area) / upPad_area;

  // canvas->Update();
  c->cd(1);

  // TPad *pull_pad = new TPad(TString::Format("%s_canv_pull", fitSummary.fit_histogram->GetName()), "Pull Pad", 0., 0., 1., 0.3);
  // pull_pad->Draw();
  // pull_pad->SetTopMargin(0.001);
  // pull_pad->SetBottomMargin(0.3);
  // pull_pad->SetGrid();

  // TPad *fit_pad = new TPad(TString::Format("%s_canv_fit", fitSummary.fit_histogram->GetName()), "Fit Pad", 0., 0.3, 1., 1.);
  // fit_pad->Draw();
  // fit_pad->SetBottomMargin(0.001);
  // fit_pad->cd();

  auto fit_Xaxis = fitSummary.fit_histogram->GetXaxis();
  auto fit_Yaxis = fitSummary.fit_histogram->GetYaxis();
  fit_Xaxis->SetLabelOffset(0.02);
  fit_Xaxis->SetLabelSize(fit_label_size);
  fit_Xaxis->SetTitleOffset(1.4);
  fit_Xaxis->SetTitleSize(fit_title_size);

  fit_Yaxis->SetLabelOffset(0.01);
  fit_Yaxis->SetLabelSize(fit_label_size);
  fit_Yaxis->SetTitleSize(fit_title_size);
  fit_Yaxis->SetTitleOffset(0);

  fitSummary.fit_histogram->Draw();

  TLegend *leg1 = new TLegend(0.55, 0.45, 0.87, 0.89);
  leg1->SetFillColor(kWhite);
  leg1->SetLineColor(kWhite);

  for (const auto &[name, title] : fitSummary.pdf_names)
  {
    leg1->AddEntry(fitSummary.fit_histogram->findObject(name.c_str()), title.c_str(), "LP");
  }

  for (const auto &par : fitSummary.fitted_params)
  {
    if (!par.constant)
    {
      leg1->AddEntry((TObject *)0, par.ToString().c_str(), "");
    }
  }

  leg1->Draw();

  c->cd(2);

  Double_t pull_min = fitSummary.pull_histogram->GetMinimum();
  Double_t pull_max = fitSummary.pull_histogram->GetMaximum();

  // std::cout << "### (" << fitSummary.pull_histogram->GetName() << ") PULL MIN: " << pull_min << " PULL MAX: " << pull_max << std::endl;

  double bound = 0;
  if (TMath::Abs(pull_min) > TMath::Abs(pull_max))
  {
    bound = TMath::Abs(pull_min);
  }
  else
  {
    bound = TMath::Abs(pull_max);
  }

  auto pull_Xaxis = fitSummary.pull_histogram->GetXaxis();
  auto pull_Yaxis = fitSummary.pull_histogram->GetYaxis();
  pull_Xaxis->SetLabelOffset(0.02);
  pull_Xaxis->SetLabelSize(pull_label_size);
  pull_Xaxis->SetTitleSize(pull_title_size);
  pull_Xaxis->SetTitle("");

  pull_Yaxis->SetLabelOffset(0.01);
  pull_Yaxis->SetLabelSize(pull_label_size);
  pull_Yaxis->SetTitleSize(pull_title_size);
  pull_Yaxis->SetTitleOffset(0.45);
  pull_Yaxis->SetTitle("Pull Distribution");
  fitSummary.pull_histogram->SetTitle("");

  fitSummary.pull_histogram->SetMaximum(bound);
  fitSummary.pull_histogram->SetMinimum(-bound);

  fitSummary.pull_histogram->Draw();

  auto line_mid = new TLine(B_MASS_VAR_MIN, 0, B_MASS_VAR_MAX, 0);
  line_mid->SetLineColor(kOrange + 4);
  line_mid->Draw();

  auto line_up = new TLine(B_MASS_VAR_MIN, bound / 2, B_MASS_VAR_MAX, bound / 2);
  line_up->SetLineStyle(kDashed);
  line_up->SetLineColor(kOrange + 4);
  line_up->Draw();

  auto line_low = new TLine(B_MASS_VAR_MIN, -bound / 2, B_MASS_VAR_MAX, -bound / 2);
  line_low->SetLineStyle(kDashed);
  line_low->SetLineColor(kOrange + 4);
  line_low->Draw();

  c->Draw();
  c->SaveAs(TString::Format("output_files/analysis/%s/%s.pdf", folder, name.Data()).Data());
}

void PrintProgress(const char *title, unsigned int total, unsigned int every, unsigned int current)
{
  if ((current + 1) % every == 0 || current + 1 == total)
  {
    std::cout << "["
              << title
              << "] Processed event: " << current + 1 << " (" << TString::Format("%.2f", ((double)(current + 1) / (double)total) * 100.) << "%)" << std::endl;
  }
}

RooPlot *CreateRooFitHistogram(TH1D *hist)
{
  RooRealVar roo_var_mass("var_mass", "B Mass Variable", B_MASS_VAR_MIN, B_MASS_VAR_MAX);
  roo_var_mass.setRange("fitting_range", B_MASS_VAR_MIN, B_MASS_VAR_MAX);

  RooDataHist roohist_B_M("roohist_B_M", "B Mass Histogram", roo_var_mass, RooFit::Import(*hist));

  RooPlot *roo_frame_mass = roo_var_mass.frame(RooFit::Title(hist->GetTitle()), RooFit::Name(TString::Format("%s_rplt", hist->GetName())));
  roohist_B_M.plotOn(roo_frame_mass, RooFit::Binning(B_MASS_HIST_BINS), RooFit::Name("B Mass Distribution"));
  roo_frame_mass->GetXaxis()->SetTitle(hist->GetXaxis()->GetTitle());

  return roo_frame_mass;
}

RooFitSummary CreateRooDataSetAndFitCB(TTree *dataSet, TString var_name, TString xAxis, bool hasExpBkg, bool useExtShape, ShapeParamters extShape, Double_t fit_low = 0, Double_t fit_up = 0)
{
  auto suffix_name = [name = dataSet->GetName()](const char *text)
  {
    return TString::Format("%s_%s", text, name);
  };

  Double_t fitRangeUp = fit_low == 0 ? B_MASS_VAR_MIN : fit_low;
  Double_t fitRangeLow = fit_up == 0 ? B_MASS_VAR_MAX : fit_up;
  
  RooRealVar roo_var_mass(var_name, "B Mass Variable", B_MASS_VAR_MIN, B_MASS_VAR_MAX);
  const char *fitting_range_name = "fitting_range";
  roo_var_mass.setRange(fitting_range_name, fitRangeUp, fitRangeLow);

  TString dataset_name = suffix_name("roodataset_B_M");
  //RooDataHist roodataset_B_M(hist_name, "B Mass Histogram", roo_var_mass, RooFit::Import(*hist));
  RooDataSet roodataset_B_M(dataset_name, "B Mass Data Set", roo_var_mass, RooFit::Import(*dataSet));

  RooPlot *roo_frame_mass = roo_var_mass.frame(RooFit::Title(dataSet->GetTitle()), RooFit::Name(TString::Format("%s_rplt", dataSet->GetName())));
  roodataset_B_M.plotOn(roo_frame_mass, RooFit::Name(dataset_name));
  roo_frame_mass->GetXaxis()->SetTitle(xAxis);

  // Crystal Ball for Signal
  RooRealVar var_mass_x0(suffix_name("var_mass_x0"), "#mu", 5278., 5170., 5500.);
  RooRealVar var_mass_sigmaLR(suffix_name("var_mass_sigmaLR"), "#sigma_{LR}", 16., 5., 40.);

  RooRealVar var_mass_alphaL(suffix_name("var_mass_alphaL"), "#alpha_{L}", 2., 0., 4.);
  RooRealVar var_mass_nL(suffix_name("var_mass_nL"), "n_{L}", 5., 0., 15.);
  RooRealVar var_mass_alphaR(suffix_name("var_mass_alphaR"), "#alpha_{R}", 2., 0., 4.);
  RooRealVar var_mass_nR(suffix_name("var_mass_nR"), "n_{R}", 5., 0., 15.);

  if (useExtShape)
  {
    var_mass_alphaL.setConstant(true);
    var_mass_alphaL.setVal(extShape.alpha_left);

    var_mass_nL.setConstant(true);
    var_mass_nL.setVal(extShape.n_left);

    var_mass_alphaR.setConstant(true);
    var_mass_alphaR.setVal(extShape.alpha_right);

    var_mass_nR.setConstant(true);
    var_mass_nR.setVal(extShape.n_right);
  }

  TString signal_name = suffix_name("sig_cb");
  RooCrystalBall sig_cb(signal_name, "CB Signal", roo_var_mass, var_mass_x0, var_mass_sigmaLR, var_mass_alphaL, var_mass_nL, var_mass_alphaR, var_mass_nR);

  std::vector<FittedParam> fitted_params{};
  std::map<std::string, std::string> pdf_names{};

  pdf_names[signal_name.Data()] = sig_cb.getTitle().Data();

  TString pull_compare_name{};
  if (hasExpBkg)
  {
    // Exponential for Background
    RooRealVar var_mass_bkg_c(suffix_name("var_mass_bkg_c"), "#lambda", -0.0014, -0.004, -0.000);

    TString background_name = suffix_name("bgk_exp");
    RooExponential bkg_exp(background_name, "Exp Background", roo_var_mass, var_mass_bkg_c);
    pdf_names[background_name.Data()] = bkg_exp.getTitle().Data();

    RooRealVar var_mass_nsig(suffix_name("nsig"), "N_{Sig}", 0., dataSet->GetEntries());
    RooRealVar var_mass_nbkg(suffix_name("nbkg"), "N_{Bkg}", 0., dataSet->GetEntries());

    TString fitted_pdf_name = suffix_name("sigplusbkg");
    RooAddPdf fitted_pdf(fitted_pdf_name, "Sig + Bkg", RooArgList(sig_cb, bkg_exp), RooArgList(var_mass_nsig, var_mass_nbkg));
    pdf_names[fitted_pdf_name.Data()] = fitted_pdf.getTitle().Data();

    RooFitResult *fitres = fitted_pdf.fitTo(roodataset_B_M, RooFit::Save(), RooFit::PrintLevel(1), RooFit::Range(fitting_range_name));

    // sigplusbkg.plotOn(roo_frame_mass, RooFit::VisualizeError(*fitres, 1), RooFit::FillColor(kOrange + 1), RooFit::FillStyle(3144));
    fitted_pdf.plotOn(roo_frame_mass, RooFit::LineColor(kRed), RooFit::LineStyle(kSolid), RooFit::Range(fitting_range_name), RooFit::Name(fitted_pdf_name));
    fitted_pdf.plotOn(roo_frame_mass, RooFit::Components(RooArgSet(bkg_exp)), RooFit::LineColor(kBlue - 7), RooFit::LineStyle(kDashed), RooFit::Range(fitting_range_name), RooFit::Name(background_name));
    fitted_pdf.plotOn(roo_frame_mass, RooFit::Components(RooArgSet(sig_cb)), RooFit::LineColor(kRed - 7), RooFit::LineStyle(kDashed), RooFit::Range(fitting_range_name), RooFit::Name(signal_name));

    fitted_params.push_back(FittedParam(var_mass_nsig, 2));
    fitted_params.push_back(FittedParam(var_mass_nbkg, 2));
    

    double sig_val = var_mass_nsig.getVal();
    double sig_err = var_mass_nsig.getError();
    double bkg_val = var_mass_nbkg.getVal();
    double bkg_err = var_mass_nbkg.getError();

    double sig_over_bkg_val = sig_val / TMath::Sqrt(sig_val + bkg_val);
    double err_prop_sig = (sig_val + 2 * bkg_val)/(2 * TMath::Power((sig_val + bkg_val), (3/2)));
    double err_prop_bkg = -sig_val/(2 * TMath::Power((sig_val + bkg_val), (3/2)));
    double sig_over_bkg_err = TMath::Sqrt(TMath::Sq(err_prop_sig * sig_err) + TMath::Sq(err_prop_bkg * bkg_err));

    fitted_params.push_back(FittedParam("sig_over_bkg", "N_{Sig}/#sqrt{N_{Sig} + N_{Bkg}}", sig_over_bkg_val, sig_over_bkg_err, 2));

    fitted_params.push_back(FittedParam(var_mass_bkg_c, 5));

    pull_compare_name = fitted_pdf_name;
  }
  else
  {
    RooFitResult *fitres = sig_cb.fitTo(roodataset_B_M, RooFit::Save(), RooFit::PrintLevel(1), RooFit::Range(fitting_range_name));
    // sigplusbkg.plotOn(roo_frame_mass, RooFit::VisualizeError(*fitres, 1), RooFit::FillColor(kOrange + 1), RooFit::FillStyle(3144));
    sig_cb.plotOn(roo_frame_mass, RooFit::LineColor(kRed), RooFit::LineStyle(kSolid), RooFit::Range(fitting_range_name), RooFit::Name(signal_name));
    pull_compare_name = signal_name;
  }

  fitted_params.push_back(FittedParam(var_mass_x0, 2));
  fitted_params.push_back(FittedParam(var_mass_sigmaLR, 2));
  fitted_params.push_back(FittedParam(var_mass_alphaL, 2));
  fitted_params.push_back(FittedParam(var_mass_nL, 2));
  fitted_params.push_back(FittedParam(var_mass_alphaR, 2));
  fitted_params.push_back(FittedParam(var_mass_nR, 2));

  RooPlot *roo_frame_pull = roo_var_mass.frame(RooFit::Title("Pull Distribution"), RooFit::Name(TString::Format("%s_pull_rplt", dataSet->GetName())));
  RooHist *pull_hist = roo_frame_mass->pullHist(dataset_name, pull_compare_name);
  roo_frame_pull->addPlotable(pull_hist, "P");

  // Double_t pull_min = pull_hist->GetMinimum();
  // Double_t pull_max = pull_hist->GetMaximum();

  // std::cout << "##### (" << pull_hist->GetName() << ") PULL MIN: " << pull_min << " PULL MAX: " << pull_max << std::endl;

  return RooFitSummary{
      roo_frame_mass,
      roo_frame_pull,
      pdf_names,
      fitted_params,
      ShapeParamters{
          var_mass_alphaL.getVal(),
          var_mass_nL.getVal(),
          var_mass_alphaR.getVal(),
          var_mass_nR.getVal()}};
}

bool InRange(double value, double center, double low_intvl, double up_intvl)
{
  return center - low_intvl < value && value < center + up_intvl;
}

bool InRange(double value, double center, double intvl)
{
  return InRange(value, center, intvl, intvl);
}

#endif