// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "CellPopulation.h"

#include <algorithm>
#include <cmath>
#include <limits>

std::pair<double, double> GivenCellPopulation(const std::vector<double> &relative_volume) {
    constexpr double everything = std::numeric_limits<double>::infinity();
    // Below this many frames a histogram has no valleys to speak of.
    constexpr size_t MIN_FRAMES = 50;
    // A valley separates two populations only where it dips below this fraction of the kept peak; a
    // shallower dip is noise in one population's tail.
    constexpr double VALLEY_FRACTION = 0.5;
    // A population holding less than this share of the frames is a stray group, not a crystal form.
    constexpr double MIN_SHARE = 0.1;

    const size_t n = relative_volume.size();
    if (n < MIN_FRAMES)
        return {-everything, everything};

    std::vector<double> v = relative_volume;
    std::sort(v.begin(), v.end());
    // The bin: twice a robust standard deviation over n^(1/3) (Freedman-Diaconis with the median
    // absolute deviation in place of the interquartile range, which a second population inflates).
    const double median = v[n / 2];
    std::vector<double> dev(n);
    for (size_t i = 0; i < n; i++)
        dev[i] = std::fabs(v[i] - median);
    std::nth_element(dev.begin(), dev.begin() + n / 2, dev.end());
    const double bin = 2.0 * 1.4826 * dev[n / 2] / std::cbrt(static_cast<double>(n));
    if (!(bin > 0))
        return {-everything, everything};

    const double lo = v.front();
    const int nbins = static_cast<int>((v.back() - lo) / bin) + 1;
    std::vector<double> count(nbins, 0.0);
    for (double x : v)
        count[std::min(nbins - 1, static_cast<int>((x - lo) / bin))] += 1.0;
    // Three-bin running mean, so a single empty or full bin is not a valley or a peak.
    std::vector<double> smooth(nbins, 0.0);
    for (int i = 0; i < nbins; i++) {
        const int a = std::max(0, i - 1), b = std::min(nbins - 1, i + 1);
        double s = 0;
        for (int j = a; j <= b; j++)
            s += count[j];
        smooth[i] = s / (b - a + 1);
    }

    // Out from a peak on each side: the valley is the lowest point, below VALLEY_FRACTION of the peak,
    // before the histogram rises again to twice its height by more than counting noise - the next
    // population, however small. None: the population runs out to the end of the data.
    const auto valley = [&](int peak, int step) {
        int lowest = peak;
        for (int i = peak + step; i >= 0 && i < nbins; i += step) {
            if (smooth[i] < smooth[lowest])
                lowest = i;
            const double low = smooth[lowest], here = smooth[i];
            // A three-bin mean of Poisson counts has a variance of a third of its value.
            const bool rises = here >= 2.0 * low && here - low > 3.0 * std::sqrt((here + low) / 3.0);
            if (low < VALLEY_FRACTION * smooth[peak] && rises)
                return lowest;
        }
        return -1;
    };

    // Every population: a local maximum of the histogram and the valleys about it. Of those that hold
    // a real share of the frames, the one centred nearest the given cell is kept - not merely the peak
    // nearest it, which can be a few stray frames when the given cell is a percent off the crystal.
    double best_low = -everything, best_high = everything, best_distance = everything;
    for (int i = 0; i < nbins; i++) {
        if (smooth[i] <= 0 || (i > 0 && smooth[i - 1] > smooth[i]) || (i + 1 < nbins && smooth[i + 1] >= smooth[i]))
            continue;
        const int left = valley(i, -1), right = valley(i, +1);
        const double low = left < 0 ? -everything : lo + (left + 0.5) * bin;
        const double high = right < 0 ? everything : lo + (right + 0.5) * bin;
        size_t count = 0;
        double sum = 0;
        for (double x : v)
            if (x >= low && x <= high) {
                count++;
                sum += x;
            }
        if (count < MIN_SHARE * static_cast<double>(n))
            continue;
        const double distance = std::fabs(sum / static_cast<double>(count));
        if (distance < best_distance) {
            best_distance = distance;
            best_low = low;
            best_high = high;
        }
    }
    return {best_low, best_high};
}