#pragma once

#ifndef clip_h
#define clip_h

#include <cmath>

#define INF 1e9

class SigmaClip {
private:
    int maxiters;
    double sigma, minValue, maxValue;
    double (*cenF)(double *, int) = nullptr;
    double (*stdF)(double *, int) = nullptr;

public:
    SigmaClip(double sigma = 3, int maxiters = 5, double (*cenF)(double *, int) = nullptr,
              double (*stdF)(double *, int) = nullptr);
    ~SigmaClip();
    double *clip(double *data, int L);

private:
    void computeBound(double *data, int L);
};

double GetAvg(double *data, int L) {
    double sum = 0;
    for (int i = 0; i < L; i++) sum += data[i];
    return sum / L;
}

double GetStd(double *data, int L) {
    double sum = 0;
    double mean = GetAvg(data, L);
    for (int i = 0; i < L; i++) sum += (data[i] - mean) * (data[i] - mean);
    return sqrt(sum / L);
}

SigmaClip::SigmaClip(double sigma, int maxiters, double (*cenF)(double *, int), double (*stdF)(double *, int)) {
    this->sigma = sigma;
    this->maxiters = maxiters;
    this->minValue = INF;
    this->maxValue = -INF;

    this->cenF = cenF == nullptr ? GetAvg : cenF;
    this->stdF = stdF == nullptr ? GetStd : stdF;
}

SigmaClip::~SigmaClip() {}

void SigmaClip::computeBound(double *data, int L) {
    double std = (*stdF)(data, L);
    double mean = (*cenF)(data, L);

    minValue = mean - (std * sigma);
    maxValue = mean + (std * sigma);
}

double *SigmaClip::clip(double *data, int L) {
    int tmp, cnt;
    double *value;

    for (int i = 1; i <= maxiters; i++) {
        computeBound(data, L);

        cnt = 0;
        for (int j = 0; j < L; j++)
            if (data[j] >= minValue && data[j] <= maxValue) cnt += 1;

        tmp = L, L = cnt, cnt = 0;
        value = new double[L];
        for (int j = 0; j < tmp; j++)
            if (data[j] >= minValue && data[j] <= maxValue) {
                value[cnt] = data[j];
                cnt += 1;
            }

        delete data;
        data = new double[L];
        for (int j = 0; j < L; j++) data[j] = value[j];
        delete[] value;
    }

    value = new double[L + 1];
    value[0] = L;
    for (int j = 0; j < L; j++) value[j + 1] = data[j];

    return value;
}

#endif