// Thiago de Paula Oliveira
#include <RcppArmadillo.h>
#include <cmath>
#include <algorithm>
using namespace Rcpp;

// [[Rcpp::plugins(openmp)]]
// [[Rcpp::depends(RcppArmadillo)]]

#include "matrixCorr_detail.h"
#include "matrixCorr_omp.h"
using matrixCorr_detail::moments::sample_var;

inline double mean_ba_ptr(const double* x, int n) {
  if (n <= 0) return NA_REAL;
  long double s = 0.0L;
  for (int i = 0; i < n; ++i) s += x[i];
  return static_cast<double>(s / static_cast<long double>(n));
}

// [[Rcpp::export]]
List bland_altman_cpp(NumericVector group1,
                      NumericVector group2,
                      double loa_multiplier = 1.96,
                      int    mode       = 1,
                      double conf_level = 0.95,
                      int n_threads = 1,
                      bool include_legacy_payload = true) {
  if (group1.size() != group2.size())
    stop("Error in bland.altman.stats: groups differ in length.");
  if (loa_multiplier <= 0.0)
    stop("Error in bland.altman.stats: improper value of 'loa_multiplier'.");
  if (mode != 1 && mode != 2)
    stop("Error in bland.altman.stats: mode must be either 1 or 2.");
  if (!(conf_level > 0.0 && conf_level < 1.0))
    stop("conf_level must be in (0,1).");
  omp_set_num_threads(std::max(1, n_threads));

  const int n_all = group1.size();
  // -- pairwise NA removal -----------------------------------------------------
  std::vector<double> g1;
  std::vector<double> g2;
  g1.reserve(n_all);
  g2.reserve(n_all);
  for (int i = 0; i < n_all; ++i) {
    const double a = group1[i];
    const double b = group2[i];
    if (!NumericVector::is_na(a) && !NumericVector::is_na(b)) {
      g1.push_back(a);
      g2.push_back(b);
    }
  }
  const int n = static_cast<int>(g1.size());
  if (n < 2) {
    warning("Warning in bland.altman.stats: less than 2 data pairs after deleting NAs.");
  }

  // -- compute means and diffs -------------------------------------------------
  NumericVector means(n), diffs(n);
  for (int i = 0; i < n; ++i) {
    const double a = g1[i];
    const double b = g2[i];
    means[i] = (a + b) / 2.0;
    diffs[i] = (mode == 1) ? (a - b) : (b - a);
  }

  // -- summary stats -----------------------------------------------------------
  const double mean_diffs = mean_ba_ptr(diffs.begin(), n);

  double sd_diffs = NA_REAL;
  if (n >= 2) {
    arma::vec diffs_view(diffs.begin(), static_cast<arma::uword>(n), /*copy_aux_mem*/ false, /*strict*/ true);
    sd_diffs = std::sqrt(sample_var(diffs_view));
  }

  const double critical   = loa_multiplier * sd_diffs;
  const double lower      = mean_diffs - critical;
  const double upper      = mean_diffs + critical;

  // -- t quantiles for CIs ------------------------------
  const double t1 = R::qt((1.0 - conf_level) / 2.0, n - 1, /*lower_tail*/1, /*log_p*/0);
  const double t2 = R::qt((conf_level + 1.0) / 2.0,       n - 1, /*lower_tail*/1, /*log_p*/0);

  // SEs per Bland & Altman 1986
  const double se_limit = sd_diffs * std::sqrt(3.0 / n);
  const double se_mean  = sd_diffs / std::sqrt(static_cast<double>(n));

  NumericVector lines = NumericVector::create(
    _["lower.limit"] = lower,
    _["mean.diffs"]  = mean_diffs,
    _["upper.limit"] = upper
  );

  NumericVector CI_lines = NumericVector::create(
    _["lower.limit.ci.lower"] = lower + t1 * se_limit,
    _["lower.limit.ci.upper"] = lower + t2 * se_limit,
    _["mean.diff.ci.lower"]   = mean_diffs + t1 * se_mean,
    _["mean.diff.ci.upper"]   = mean_diffs + t2 * se_mean,
    _["upper.limit.ci.lower"] = upper + t1 * se_limit,
    _["upper.limit.ci.upper"] = upper + t2 * se_limit
  );

  if (include_legacy_payload) {
    DataFrame groups = DataFrame::create(
      _["group1"] = wrap(g1),
      _["group2"] = wrap(g2),
      _["stringsAsFactors"] = false
    );

    return List::create(
      _["means"]         = means,
      _["diffs"]         = diffs,
      _["groups"]        = groups,
      _["based.on"]      = n,
      _["lower.limit"]   = lower,
      _["mean.diffs"]    = mean_diffs,
      _["upper.limit"]   = upper,
      _["lines"]         = lines,
      _["CI.lines"]      = CI_lines,
      _["loa_multiplier"] = loa_multiplier,
      _["critical.diff"] = critical
    );
  }

  return List::create(
    _["means"]         = means,
    _["diffs"]         = diffs,
    _["based.on"]      = n,
    _["lower.limit"]   = lower,
    _["mean.diffs"]    = mean_diffs,
    _["upper.limit"]   = upper,
    _["lines"]         = lines,
    _["CI.lines"]      = CI_lines,
    _["loa_multiplier"] = loa_multiplier,
    _["critical.diff"] = critical
  );
}

// [[Rcpp::export]]
int ba_openmp_threads() {
  int n = 1;
#ifdef _OPENMP
  n = omp_get_max_threads();
#endif
  return n;
}
