// fit_WL.cxx
// (c) 2026 Ako_b
// see for details
// https://root-forum.cern.ch/t/question-about-the-scaling-factor-in-root-s-weighted-poisson-log-likelihood-wls-fit/64952/11

// statistical model
static constexpr int g_npar = 1; // number of model parameters
double model_prediction_for_bin(const double *x, const double *p) {
    double mean_number_of_events_in_bin = p[0];

    return mean_number_of_events_in_bin;
}

struct CustomNLL {
    const int mode = 2; //

    const ROOT::Fit::BinData fData;

    CustomNLL(int mode, const ROOT::Fit::BinData& data) : mode{mode}, fData(data) {}

    // fitter expects an operator() that takes a pointer to the parameter array
    double operator()(const double* p) const {
        double nll = 0.0;

        // Loop over all bins in the BinData structure
        for (unsigned int i = 0; i < fData.Size(); ++i) {
            double x = fData.Coords(i)[0]; // Bin center
            double y = fData.Value(i);     // Observed content (counts)
            double err = fData.Error(i);
            if (y <= 0.0 || err <= 0.0) continue;

            double mu = model_prediction_for_bin(&x, p);
            mu = std::max(mu, 0.0);

            const double s = (err * err) / y; // Bin scale factor s_i = sigma_i^2 / y_i
            double weight = 1.0;

            if (mode == 1) {
                weight = s;       // Variant 1: s_i
            } else if (mode == 2) {
                weight = 1.0;     // Variant 2: 1.0
            } else if (mode == 3) {
                weight = 1.0 / s; // Variant 3: 1/s_i
            } else {
                std::cerr << "Unrecognized mode " << mode << std::endl;
            }

            // Deviance term: mu - y - y * log(mu / y)
            nll += weight * (mu - y - y * ROOT::Math::Util::EvalLog(mu / y));
        }
        return 2.0 * nll;
    }
};

double FillHisto(TH1D& h_toy, double mean_number_of_events_in_bin_true, double scale_factor_true) {
    h_toy.Reset();

    double n_observed;
    double n_eff;

    double true_parameters_arr[1] = { mean_number_of_events_in_bin_true };

    // Fill first bin
    double x1[1] = { h_toy.GetBinCenter(1) };
    double first_bin_mean = model_prediction_for_bin(x1, true_parameters_arr);

    n_observed = gRandom->Poisson(first_bin_mean);
    h_toy.SetBinContent(1, n_observed);
    h_toy.SetBinError(1, sqrt(n_observed));

    n_eff += n_observed;

    // Fill second bin with different exposure
    double x2[1] = { h_toy.GetBinCenter(2) };
    double second_bin_mean = model_prediction_for_bin(x2, true_parameters_arr);

    n_observed = gRandom->Poisson(second_bin_mean / scale_factor_true);
    h_toy.SetBinContent(2, n_observed * scale_factor_true);
    h_toy.SetBinError(2, sqrt(n_observed) * scale_factor_true);

    n_eff += n_observed;

    return n_eff;
}



std::optional<double> CustomFit(int mode, const TH1D &h_toy, ROOT::Fit::Fitter &fitter) {
    ROOT::Fit::DataOptions options;
    ROOT::Fit::BinData bin_data(options);
    ROOT::Fit::FillData(bin_data, &h_toy);

    CustomNLL nll(mode, bin_data);

    std::vector<double> initial_param_grid = { 10.0, 1000.0 };
    size_t grid_size = initial_param_grid.size();

    std::vector<ROOT::Fit::FitResult> results;

    for (size_t i = 0; i < grid_size; ++i) {
        double initial_param_value = initial_param_grid[i];

        double initial_params[g_npar] = { initial_param_value };
        bool fit_ok = fitter.FitFCN(g_npar, nll, initial_params, bin_data.Size());

        if (fit_ok && fitter.Result().IsValid()) {
            results.push_back(fitter.Result());
        }
    }

    if (! results.empty()) {
        std::vector<ROOT::Fit::FitResult>::iterator it_minFCN;

        it_minFCN = std::min_element(results.begin(), results.end(), [](const ROOT::Fit::FitResult &a, const ROOT::Fit::FitResult &b)
        {
            return a.MinFcnValue() < b.MinFcnValue();
        });

        return (*it_minFCN).Parameter(0);
    }

    return std::nullopt;
}

void drawLegend(double x1, double y1, double x2, double y2, std::vector<TH1*> histogram_pointers) {
    TLegend legend(x1, y1, x2, y2);

    for (TH1* h: histogram_pointers) {
        TString title = h->GetTitle();
        legend.AddEntry(h, title, "l");
    }
    legend.SetTextSize(0.03);

    legend.DrawClone();
}

void drawCustomStatBox(double x1, double y1, double x2, double y2, std::vector<TH1*> histogram_pointers, double mean_number_of_events_in_bin_true, double scale_factor_true) {
    TPaveText pt(x1, y1, x2, y2, "NDC");

    // Define grid spacing
    int numRows = histogram_pointers.size() + 1; // +1 for the header row
    double rowHeight = 1.0 / numRows;

    // column x coordinates 
    double col1_x = 0.02; // Left column
    double col2_x = 0.42; // Middle column
    double col3_x  = 0.70; // Right column

    double current_row_y = 1.0 - (rowHeight * 0.5); // Center line of the row

    pt.SetFillColor(0); // white
    pt.SetTextAlign(12); // center align
    pt.SetTextSize(0.03);

    // draw header
    pt.SetTextFont(62); // bold
    pt.AddText(col1_x, current_row_y, "Histogram");
    pt.AddText(col2_x, current_row_y, "bias");
    pt.AddText(col3_x,  current_row_y, "overdispersion");

    // draw horizontal line under header
    pt.AddLine(0.0, 1.0 - rowHeight, 1.0, 1.0 - rowHeight);

    // draw statistics rows
    pt.SetTextFont(42); // regular

    // ranged for loop preserves order
    for (TH1* h: histogram_pointers) {
        current_row_y -= rowHeight; // Move to the next row coordinate

        // Format names and numeric values
        TString name = h->GetTitle();

        double bias = h->GetMean() - mean_number_of_events_in_bin_true;
        TString bias_string = TString::Format("%+.5f", bias);

        double true_variance = mean_number_of_events_in_bin_true * (1.0 * scale_factor_true) / (1.0 + scale_factor_true);
        double true_sigma = sqrt(true_variance);
        double overdispersion = h->GetRMS() / true_sigma;
        TString overdispersion_string  = TString::Format("%.5f", overdispersion);

        std::cout << h->GetTitle() << ": bias = " << bias << ", overdispersion = " << overdispersion << std::endl; 

        pt.AddText(col1_x, current_row_y, name);
        pt.AddText(col2_x, current_row_y, bias_string);
        pt.AddText(col3_x,  current_row_y, overdispersion_string);
    }

    pt.DrawClone();
}

void fit_WL() {
    gRandom->SetSeed(2026);

    const int n_toys = 10000; // number of toy experiments
    const double mean_number_of_events_in_bin_true = 100;
    const double scale_factor_true = 1.0 / 100.0;

    int nbins = 1000;
    const double h_toy_xlow = mean_number_of_events_in_bin_true - 5.0 * sqrt(mean_number_of_events_in_bin_true);
    const double h_toy_xhigh = mean_number_of_events_in_bin_true + 5.0 * sqrt(mean_number_of_events_in_bin_true);


    TH1D h_WL("h_WL", "TH1::Fit, WL;#hat{#mu};Toys/bin", nbins, h_toy_xlow, h_toy_xhigh);
    TH1D h_s("h_s", "Custom fit: s#times#Deltal;#hat{#mu};Toys/bin", nbins, h_toy_xlow, h_toy_xhigh);
    TH1D h_1("h_1", "Custom fit: 1#times#Deltal;#hat{#mu};Toys/bin", nbins, h_toy_xlow, h_toy_xhigh);
    TH1D h_1_over_s("h_1_over_s", "Custom fit: 1/s#times#Deltal;#hat{#mu};Toys/bin", nbins, h_toy_xlow, h_toy_xhigh);

    h_WL.SetLineColorAlpha(kBlack, 0.8);
    h_s.SetLineColor(kRed);
    h_1.SetLineColor(kBlue);
    h_1_over_s.SetLineColor(kGreen-1);

    h_WL.SetLineStyle(kSolid);
    h_s.SetLineStyle(kDashed);
    h_1.SetLineStyle(kDashed);
    h_1_over_s.SetLineStyle(kDashed);

    h_WL.SetLineWidth(3);
    h_s.SetLineWidth(3);
    h_1.SetLineWidth(3);
    h_1_over_s.SetLineWidth(3);

    auto *f_model = new TF1("f_model", model_prediction_for_bin, 0.0, 100.0, g_npar);
    f_model->SetParNames("mean_number_of_events_in_bin");

    TH1D h_toy("h_toy", "Toy Histogram", 2, 0.0, 1.0); // two bins - the first unweighted, the second weighted

    ROOT::Fit::Fitter fitter;

    for (int toy = 0; toy < n_toys; ++toy) {
        double n_eff = FillHisto(h_toy, mean_number_of_events_in_bin_true, scale_factor_true);

        // 1. ROOT TH1::Fit with option "WL"
        f_model->SetParameters(1.0);
        h_toy.Fit(f_model, "WL Q 0");
        h_WL.Fill(f_model->GetParameter(0));

        // 2. s * Delta l
        auto res2 = CustomFit(1, h_toy, fitter);
        if (res2.has_value()) {
            h_s.Fill(res2.value());
        }

        // 3. 1 * Delta l
        auto res3 = CustomFit(2, h_toy, fitter);
        if (res3.has_value()) {
            h_1.Fill(res3.value());
        }

        // 4. 1/s * Delta l
        auto res4 = CustomFit(3, h_toy, fitter);
        if (res4.has_value()) {
            h_1_over_s.Fill(res4.value());
        }
    }

    TCanvas *c = new TCanvas("c", "Custom fit comparison", 1280, 720);

    gStyle->SetOptStat(0);

    auto *h_WL_handle = (TH1D*)h_WL.DrawCopy();
    auto *h_s_handle = (TH1D*)h_s.DrawCopy("same");
    auto *h_1_handle = (TH1D*)h_1.DrawCopy("same");
    auto *h_1_over_s_handle = (TH1D*)h_1_over_s.DrawCopy("same");

    std::vector<TH1*> histogram_pointers = {h_WL_handle, h_s_handle, h_1_handle, h_1_over_s_handle};

    // set y axis scale
    ROOT::RVecD max_values;
    for (const auto *h: histogram_pointers) {
        max_values.push_back(h->GetMaximum());
    }
    h_WL_handle->GetYaxis()->SetRangeUser(0.0, ROOT::VecOps::Max(max_values) * 1.2);

    // print number of entries to check if fit did converge for all toys
    for (const auto *h: histogram_pointers) {
        std::cout << h->GetName() << "->GetEntries() = " << h->GetEntries() << std::endl;
    }

    drawLegend(0.13, 0.65, 0.4, 0.88, histogram_pointers);
    drawCustomStatBox(0.56, 0.65, 0.90, 0.88, histogram_pointers, mean_number_of_events_in_bin_true, scale_factor_true);
}