#include "datascience/autocorr.h"

#include <cmath>

namespace fn::ds {

static double mean_inline(const double* x, std::size_t n) {
    if (n == 0) return 0.0;
    double s = 0.0;
    for (std::size_t i = 0; i < n; ++i) s += x[i];
    return s / static_cast<double>(n);
}

static double var_inline(const double* x, std::size_t n, double mu) {
    if (n == 0) return 0.0;
    double s = 0.0;
    for (std::size_t i = 0; i < n; ++i) {
        double d = x[i] - mu;
        s += d * d;
    }
    return s / static_cast<double>(n);
}

double autocorr_lag(const double* x, std::size_t n, std::size_t k) {
    if (x == nullptr || n <= 1 || k >= n) return 0.0;
    double mu = mean_inline(x, n);
    double var = var_inline(x, n, mu);
    if (var <= 0.0) return (k == 0 ? 0.0 : 0.0);

    double cov = 0.0;
    std::size_t m = n - k;
    for (std::size_t i = 0; i < m; ++i) {
        cov += (x[i] - mu) * (x[i + k] - mu);
    }
    cov /= static_cast<double>(m);
    return cov / var;
}

void autocorr_acf(const double* x, std::size_t n,
                  std::size_t max_lag, double* out) {
    if (out == nullptr || max_lag == 0) return;
    if (x == nullptr || n <= 1) {
        for (std::size_t k = 0; k < max_lag; ++k) out[k] = 0.0;
        return;
    }
    double mu = mean_inline(x, n);
    double var = var_inline(x, n, mu);
    if (var <= 0.0) {
        for (std::size_t k = 0; k < max_lag; ++k) out[k] = 0.0;
        return;
    }
    for (std::size_t k = 0; k < max_lag; ++k) {
        if (k >= n) { out[k] = 0.0; continue; }
        double cov = 0.0;
        std::size_t m = n - k;
        for (std::size_t i = 0; i < m; ++i) {
            cov += (x[i] - mu) * (x[i + k] - mu);
        }
        cov /= static_cast<double>(m);
        out[k] = cov / var;
    }
}

double autocorr_tau(const double* x, std::size_t n,
                    std::size_t max_lag, double cutoff) {
    if (x == nullptr || n <= 1) return 1.0;
    if (max_lag == 0) max_lag = 1;
    double mu = mean_inline(x, n);
    double var = var_inline(x, n, mu);
    if (var <= 0.0) return 1.0;

    double sum = 0.0;
    std::size_t kmax = (max_lag < n ? max_lag : n - 1);
    for (std::size_t k = 1; k < kmax; ++k) {
        double cov = 0.0;
        std::size_t m = n - k;
        for (std::size_t i = 0; i < m; ++i) {
            cov += (x[i] - mu) * (x[i + k] - mu);
        }
        cov /= static_cast<double>(m);
        double r = cov / var;
        if (std::fabs(r) < cutoff) break;
        sum += r;
    }
    return 1.0 + 2.0 * sum;
}

} // namespace fn::ds