#include "welford_mean.h"
#include <math.h>
#include <stdint.h>

static inline float kahanSummation(float sum_previous, float input,
                                   float *accumulator) {
  const float y = input - *accumulator;
  const float t = sum_previous + y;
  *accumulator = (t - sum_previous) - y;
  return t;
}

void welford_mean_reset(WelfordMean_t *wm) {
  if (wm) {
    wm->_mean = 0.0f;
    wm->_M2 = 0.0f;
    wm->_mean_accum = 0.0f;
    wm->_M2_accum = 0.0f;
    wm->_count = 0;
  }
}

bool welford_mean_valid(WelfordMean_t *wm) {
  if (wm == NULL) {
    return false;
  }
  return wm->_count > 2;
}

bool welford_mean_updata(WelfordMean_t *wm, float new_val) {
  if (wm == NULL) {
    return false;
  }
  if (wm->_count == 0) {
    welford_mean_reset(wm);
    wm->_count = 1;
    wm->_mean = new_val;
    return false;
  } else if (wm->_count == UINT16_MAX) {
    // count overflow
    //  reset count, but maintain mean and variance
    wm->_M2 = wm->_M2 / wm->_count;
    wm->_M2_accum = 0;
    wm->_count = 1;
  } else {
    wm->_count++;
  }

  // mean accumulates the mean of the entire dataset
  // delta can be very small compared to the mean, use algorithm to minimise
  // numerical error
  float delta = new_val - wm->_mean;
  float mean_change = delta / wm->_count;
  wm->_mean = kahanSummation(wm->_mean, mean_change, &wm->_mean_accum);

  // M2 aggregates the squared distance from the mean
  // count aggregates the number of samples seen so far
  const float M2_change = delta * (new_val - wm->_mean);
  wm->_M2 = kahanSummation(wm->_M2, M2_change, &wm->_M2_accum);

  // protect against floating point precision causing negative variances
  wm->_M2 = wm->_M2 > 0 ? wm->_M2 : 0;

  if (!isfinite(wm->_mean) || !isfinite(wm->_M2)) {
    welford_mean_reset(wm);
    return false;
  }

  return welford_mean_valid(wm);
}

int welford_mean_get_count(WelfordMean_t *wm) {
  if (wm) {
    return wm->_count;
  } else {
    return 0;
  }
}

float welford_mean_get_mean(WelfordMean_t *wm) {
  if (wm) {
    return wm->_mean;
  } else {
    return NAN;
  }
}

float welford_mean_get_variance(WelfordMean_t *wm) {
  if (wm) {
    return wm->_M2 / (wm->_count - 1);
  } else {
    return NAN;
  }
}

float welford_mean_get_stddev(WelfordMean_t *wm) {
  if (wm) {
    return sqrtf(welford_mean_get_variance(wm));
  } else {
    return NAN;
  }
}

bool welford_mean_vector3f_valid(WelfordMeanVector3f_t *wm) {
  if (wm == NULL) {
    return false;
  }
  return wm->_count > 2;
}

void welford_mean_vector3f_reset(WelfordMeanVector3f_t *wm) {
  if (wm) {
    for (int i = 0; i < 3; i++) {
      wm->_mean[i] = 0.0f;
      wm->_mean_accum[i] = 0.0f;
      for (int j = 0; j < 3; j++) {
        wm->_M2[i][j] = 0.0f;
        wm->_M2_accum[i][j] = 0.0f;
      }
    }
    wm->_count = 0;
  }
}

bool welford_mean_vector3f_updata(WelfordMeanVector3f_t *wm, float new_val[3]) {
  if (wm == NULL || new_val == NULL) {
    return false;
  }

  if (wm->_count == 0) {
    welford_mean_vector3f_reset(wm);
    wm->_count = 1;
    for (int i = 0; i < 3; i++) {
      wm->_mean[i] = new_val[i];
    }
    return false;

  } else if (wm->_count == UINT16_MAX) {
    // count overflow
    //  reset count, but maintain mean and variance
    for (int i = 0; i < 3; i++) {
      for (int j = 0; j < 3; j++) {
        wm->_M2[i][j] = wm->_M2[i][j] / wm->_count;
        wm->_M2_accum[i][j] = 0;
      }
    }

    wm->_count = 1;

  } else {
    wm->_count++;
  }

  // mean
  //  accumulates the mean of the entire dataset
  //  use Kahan summation because delta can be very small compared to the mean
  float delta[3];
  for (int i = 0; i < 3; i++) {
    delta[i] = new_val[i] - wm->_mean[i];
  }
  float y[3];
  float t[3];
  for (int i = 0; i < 3; i++) {
    y[i] = delta[i] - wm->_mean_accum[i];
    t[i] = wm->_mean[i] + y[i];
    wm->_mean_accum[i] = (t[i] - wm->_mean[i]) - y[i];
    wm->_mean[i] = t[i];
  }

  for (int i = 0; i < 3; ++i) {
    if (!isfinite(wm->_mean[i])) {
      welford_mean_vector3f_reset(wm);
      return false;
    }
  }

  // covariance
  //  Kahan summation (upper triangle only)
  // eg C(x,y) += dx * (y - mean_y)
  float m2_change[3][3];

  for (size_t r = 0; r < 3; r++) {
    for (size_t c = r; c < 3; c++) {
      m2_change[r][c] = delta[r] * (new_val[c] - wm->_mean[c]);
    }
  }

  for (size_t r = 0; r < 3; r++) {
    for (size_t c = r; c < 3; c++) {

      float y = m2_change[r][c] - wm->_M2_accum[r][c];
      float t = wm->_M2[r][c] + y;
      wm->_M2_accum[r][c] = (t - wm->_M2[r][c]) - y;

      wm->_M2[r][c] = t;
    }

    // protect against floating point precision causing negative variances
    if (wm->_M2[r][r] < 0) {
      wm->_M2[r][r] = 0;
    }
  }

  // make symmetric
  for (size_t r = 0; r < 3; r++) {
    for (size_t c = r + 1; c < 3; c++) {
      wm->_M2[c][r] = wm->_M2[r][c];
    }
  }

  for (size_t r = 0; r < 3; r++) {
    for (size_t c = 0; c < 3; c++) {
      if (!isfinite(wm->_M2[r][c])) {
        welford_mean_vector3f_reset(wm);
        return false;
      }
    }
  }

  return welford_mean_vector3f_valid(wm);
}

int welford_mean_vector3f_get_count(WelfordMeanVector3f_t *wm) {
  if (wm == NULL) {
    return 0;
  }

  return wm->_count;
}

bool welford_mean_vector3f_get_mean(WelfordMeanVector3f_t *wm, float mean[3]) {
  if (wm == NULL || mean == NULL) {
    return false;
  }
  for (int i = 0; i < 3; i++) {
    mean[i] = wm->_mean[i];
  }
  return true;
}

bool welford_mean_vector3f_get_variance(WelfordMeanVector3f_t *wm,
                                        float variance[3]) {
  if (wm == NULL || variance == NULL) {
    return false;
  }
  for (int i = 0; i < 3; i++) {
    variance[i] = wm->_M2[i][i] / (wm->_count - 1);
  }
  return true;
}

float welford_mean_vector3f_get_covariance(WelfordMeanVector3f_t *wm, int x,
                                           int y) {
  if (wm == NULL || x < 0 || x > 2 || y < 0 || y > 2) {
    return NAN;
  }
  return wm->_M2[x][y] / (wm->_count - 1);
}