// 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 <array>
#include <cmath>
#include <limits>
#include <numbers>

#include <Eigen/Dense>

std::vector<CellPopulationRange> CellPopulations(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 {};

    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 {};

    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, if it holds a real share
    // of the frames.
    std::vector<CellPopulationRange> populations;
    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);
        // A population is its highest peak: a stray frame or two in its tail is a local maximum too, and
        // with no valley of its own its range is the whole population - or, where the tail runs out
        // before any valley, everything, two forms included.
        bool highest = true;
        for (int j = left < 0 ? 0 : left; j <= (right < 0 ? nbins - 1 : right); j++)
            if (smooth[j] > smooth[i])
                highest = false;
        if (!highest)
            continue;
        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;
        populations.push_back({low, high, count, sum / static_cast<double>(count)});
    }
    return populations;
}

namespace {
    // A cell's shape, for telling forms apart: its edges as logarithms, so a percent weighs the same on every
    // edge. Not its angles: a still's refined angles carry the indexing more than the crystal - the same form
    // comes out with its monoclinic angle on one axis or another, obtuse or acute, or held at 90 degrees - and
    // those spread more than two forms' angles differ.
    Eigen::Vector3d Shape(const UnitCell &c) {
        return {std::log(c.a), std::log(c.b), std::log(c.c)};
    }
}

std::pair<double, double> GivenCellPopulation(const std::vector<UnitCell> &cells, const UnitCell &given) {
    constexpr double everything = std::numeric_limits<double>::infinity();
    // How much less certain a population's size is than its proportions, in telling which is the given cell's.
    // The size a frame refines to carries the error of the detector distance it was refined at, all edges
    // alike, on top of the crystal's - a few percent - while a population's mean proportions come within about
    // a percent of its form's. The size still counts, at a third, so that two forms of one shape are told apart.
    constexpr double SIZE_UNCERTAINTY = 3.0;

    // The proportions: the logarithms of the edges less their mean, and the cosines of the angles, unsigned (an
    // angle and its supplement describe one lattice). The frames were indexed on the given cell and refined in
    // its crystal system, so their angles are measured on its axes (put in its order, AlignedCell) - unlike the
    // de-novo cells SeparateCellForms compares, whose angles it leaves out.
    const auto proportions = [](const UnitCell &c) {
        const auto cosine = [](float deg) { return std::fabs(std::cos(deg * std::numbers::pi / 180.0)); };
        const Eigen::Vector3d s = Shape(c);
        Eigen::Matrix<double, 6, 1> p;
        p << s.array() - s.mean(), cosine(c.alpha), cosine(c.beta), cosine(c.gamma);
        return p;
    };

    const double given_volume = gemmi::UnitCell(given).volume;
    std::vector<double> rel(cells.size());
    for (size_t i = 0; i < cells.size(); i++)
        rel[i] = gemmi::UnitCell(cells[i]).volume / given_volume - 1.0;

    // The population whose mean cell is 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 (const auto &p : CellPopulations(rel)) {
        Eigen::Matrix<double, 6, 1> mean_proportions = Eigen::Matrix<double, 6, 1>::Zero();
        double mean_size = 0.0;
        for (size_t i = 0; i < cells.size(); i++)
            if (rel[i] >= p.low && rel[i] <= p.high) {
                const UnitCell c = AlignedCell(cells[i], given);
                mean_proportions += proportions(c);
                mean_size += Shape(c).mean();
            }
        mean_proportions /= static_cast<double>(p.count);
        mean_size /= static_cast<double>(p.count);
        const double size = (mean_size - Shape(given).mean()) / SIZE_UNCERTAINTY;
        const double distance = (mean_proportions - proportions(given)).squaredNorm() + size * size;
        if (distance < best_distance) {
            best_distance = distance;
            best_low = p.low;
            best_high = p.high;
        }
    }
    return {best_low, best_high};
}

UnitCell AlignedCell(const UnitCell &cell, const UnitCell &reference) {
    const std::array<float, 3> edge = {cell.a, cell.b, cell.c}, angle = {cell.alpha, cell.beta, cell.gamma};
    const Eigen::Vector3d ref = Shape(reference);
    std::array<int, 3> p = {0, 1, 2};
    UnitCell best = cell;
    double best_distance = std::numeric_limits<double>::infinity();
    do {
        const UnitCell c{edge[p[0]], edge[p[1]], edge[p[2]], angle[p[0]], angle[p[1]], angle[p[2]]};
        const double distance = (Shape(c) - ref).squaredNorm();
        if (distance < best_distance) {
            best_distance = distance;
            best = c;
        }
    } while (std::next_permutation(p.begin(), p.end()));
    return best;
}

CellForms SeparateCellForms(const std::vector<UnitCell> &cells, const UnitCell &reference) {
    CellForms out;
    out.form.assign(cells.size(), 0);
    const size_t n = cells.size();
    if (n == 0)
        return out;
    std::vector<Eigen::Vector3d> f(n);
    Eigen::Vector3d mean = Eigen::Vector3d::Zero();
    for (size_t i = 0; i < n; i++) {
        f[i] = Shape(AlignedCell(cells[i], reference));
        mean += f[i];
    }
    mean /= static_cast<double>(n);
    Eigen::Matrix3d cov = Eigen::Matrix3d::Zero();
    for (const auto &x : f)
        cov += (x - mean) * (x - mean).transpose();
    // The direction the cells spread along most: between two forms, where there are two, as the spread
    // between them adds to the spread of each.
    const Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> eig(cov);
    const Eigen::Vector3d axis = eig.eigenvectors().col(2);
    out.coordinate.resize(n);
    for (size_t i = 0; i < n; i++)
        out.coordinate[i] = (f[i] - mean).dot(axis);

    auto populations = CellPopulations(out.coordinate);
    if (populations.size() < 2)
        return out;
    std::sort(populations.begin(), populations.end(),
              [](const CellPopulationRange &x, const CellPopulationRange &y) { return x.count > y.count; });
    const auto &major = populations[0], &minor = populations[1];
    const auto in = [](const CellPopulationRange &p, double x) { return x >= p.low && x <= p.high; };
    Eigen::Vector3d mean_major = Eigen::Vector3d::Zero(), mean_minor = mean_major;
    std::vector<UnitCell> minor_cells;
    for (size_t i = 0; i < n; i++) {
        if (in(major, out.coordinate[i]))
            mean_major += f[i];
        else if (in(minor, out.coordinate[i])) {
            mean_minor += f[i];
            minor_cells.push_back(AlignedCell(cells[i], reference));
        }
    }
    mean_major /= static_cast<double>(major.count);
    mean_minor /= static_cast<double>(minor.count);
    out.separation = (mean_minor - mean_major).norm();
    if (out.separation < MIN_FORM_SEPARATION)
        return out;
    for (size_t i = 0; i < n; i++)
        if (in(minor, out.coordinate[i]))
            out.form[i] = 1;
    out.minor_frames = minor_cells.size();
    out.minor_cell = MeanUnitCell(minor_cells);
    return out;
}