#include "stats.hpp"
#include "welford.hpp"
#include <fmt/core.h>
#include <numeric>
#include <numbers>
#include <functional>

void Stats::dump(std::string msg)
{
  fmt::println("{}: sum: {}, sumsq: {}, n: {}, "
            "min: {}, max: {}, mean: {}, stddev: {}",
            msg, sum, sumsq, n, min, max, mean(),
            stddev());
}

double simpson(double a, double b, uint32_t n, const std::function<double(double)> func) {
    const double h = ( b - a ) / n;
    double sum = 0.0;
    for ( uint32_t i = 0; i < n; ++i ) {
    	const double x = a + i * h;
    	sum += ( func(x) + 4.0 * func(x + h / 2.0) + func(x + h) ) / 6.0;
    }
    return sum * h;
}


TTest internal_t_test(auto& first, auto& second) {
  double n1 = first.n;
  double n2 = second.n;

  double mean1 = first.mean();
  double mean2 = second.mean();

  double var1 = first.sample_variance();
  double var2 = second.sample_variance();

  // calculate t-stat
  double delta_mean = mean1 - mean2;
  double pooled_se = std::sqrt((var1 / n1) + (var2 / n2));
  double t_stat = delta_mean / pooled_se;

  // calculate degrees of freedom
  double num = std::pow((var1 / n1) + (var2 / n2), 2);
  double den = (std::pow(var1 / n1, 2) / (n1 - 1)) + (std::pow(var2 / n2, 2) / (n2 - 1));
  double dof = num / den;

  // welch calculation
  double temp = first.sample_variance() / first.n + second.sample_variance() / second.n;
  double welchs = (first.mean() - second.mean()) / std::sqrt(temp);

  double gamm = std::exp(
      std::lgamma(dof / 2.0) + std::lgamma(0.5)
        - std::lgamma(dof / 2.0 + 0.5));
  double b = dof / ( welchs * welchs + dof);

  auto func = [&](const double r) {
    return std::pow(r, dof / 2.0 - 1.0) / std::sqrt(1.0 - r);
  };

  double p_val = simpson(0.0, b, 10000, func) / gamm;

  return {
    .t_stat = t_stat,
    .dof = dof,
    .p_val = p_val};
}

TTest Stats::t_test(Stats& other) {
  return internal_t_test(*this, other);
}

TTest WelfordStats::t_test(WelfordStats& other) {
  return internal_t_test(*this, other);
}