#include "statistics.h"

#include <algorithm>

#include <aocommon/uvector.h>

void WinsorizedMeanAndStdDev(const ImageInterface& image, float& mean,
                             float& stddev) {
  const size_t size = image.Width() * image.Height();
  aocommon::UVector<float> data(size);
  const float* image_data = image.Data();
  for (size_t y = 0; y != image.Height(); ++y) {
    const float* row = &image_data[y * image.Stride()];
    std::copy_n(row, image.Width(), &data[y * image.Width()]);
  }
  const size_t low_index = (size_t)std::floor(0.1 * size);
  const size_t high_index = (size_t)std::ceil(0.9 * size) - 1;
  // We need this less than operator, because the normal operator
  // does not enforce a strictly ordered set, because a<b != !(b<a) in the case
  // of nans/infs.
  const auto less_than = [](const float a, const float b) -> bool {
    if (std::isfinite(a)) {
      if (std::isfinite(b))
        return a < b;
      else
        return true;
    }
    return false;
  };
  std::nth_element(data.begin(), data.begin() + low_index, data.end(),
                   less_than);
  const float low_value = data[low_index];
  std::nth_element(data.begin() + low_index, data.begin() + high_index,
                   data.end(), less_than);
  const float high_value = data[high_index];

  // Calculate mean
  mean = 0.0;
  size_t count = 0;
  for (size_t y = 0; y != image.Height(); ++y) {
    const float* row = &image_data[y * image.Stride()];
    for (size_t x = 0; x != image.Width(); ++x) {
      if (std::isfinite(row[x])) {
        if (row[x] < low_value)
          mean += low_value;
        else if (row[x] > high_value)
          mean += high_value;
        else
          mean += row[x];
        count++;
      }
    }
  }
  if (count > 0) mean /= float(count);
  // Calculate variance
  stddev = 0.0;
  count = 0;
  const float low_minus_mean_sq = (low_value - mean) * (low_value - mean);
  const float high_minus_mean_sq = (high_value - mean) * (high_value - mean);
  for (size_t y = 0; y != image.Height(); ++y) {
    const float* row = &image_data[y * image.Stride()];
    for (size_t x = 0; x != image.Width(); ++x) {
      if (std::isfinite(row[x])) {
        if (row[x] < low_value)
          stddev += low_minus_mean_sq;
        else if (row[x] > high_value)
          stddev += high_minus_mean_sq;
        else {
          const float diff = row[x] - mean;
          stddev += diff * diff;
        }
        count++;
      }
    }
  }
  if (count > 0)
    stddev = std::sqrt(1.54 * stddev / float(count));
  else
    stddev = 0.0;
}
