// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#ifndef JUNGFRAUJOCH_REGRESSION_H
#define JUNGFRAUJOCH_REGRESSION_H

#include <vector>
#include <utility>
#include <cmath>

#include "../../common/JFJochException.h"

struct RegressionResult {
    float intercept;
    float slope;
    float r_square;
};

template <class fptype = float>
RegressionResult regression(std::vector<float> &x, std::vector<float> &y, const size_t N) {
    if (N < 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Regression: minimum 3 points needed for any meaningful result");
    if (x.size() < N || y.size() < N)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Regression: mismatch between input parameter size");

    fptype x_sum = 0.0;
    fptype y_sum = 0.0;
    fptype xy_sum = 0.0;
    fptype x2_sum = 0.0;

    int64_t nelem = 0;

    for (int i = 0; i < N; i++) {
        if (std::isfinite(y[i])) {
            x_sum += x[i];
            y_sum += y[i];
            x2_sum += x[i] * x[i];
            xy_sum += x[i] * y[i];
            nelem++;
        }
    }

    if (nelem < 3)
        return {
                .intercept = 0,
                .slope = 0,
                .r_square = 0
        };

    fptype denominator = nelem * x2_sum - x_sum * x_sum;
    fptype intercept = (y_sum * x2_sum - x_sum * xy_sum) / denominator;
    fptype slope = (nelem * xy_sum - x_sum * y_sum) / denominator;

    fptype y_mean = y_sum / nelem;
    fptype ss_tot = 0.0;
    fptype ss_reg = 0.0;

    for (int i = 0; i < nelem; i++) {
        fptype pred = slope * x[i] + intercept;
        ss_tot += (y[i] - y_mean) * (y[i] - y_mean);
        ss_reg += (pred - y_mean) * (pred - y_mean);
    }

    fptype r_square = 0.0;

    if (ss_tot != 0.0)
        r_square = ss_reg / ss_tot;

    return {
            .intercept = static_cast<float>(intercept),
            .slope = static_cast<float>(slope),
            .r_square = static_cast<float>(r_square)
    };
};

#endif //JUNGFRAUJOCH_REGRESSION_H