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

#include <algorithm>
#include <array>
#include <cmath>

#include "HarmonicContamination.h"

namespace {

    // A frame's orientation as the map from a spot's reciprocal vector to its Miller indices, fitted
    // by least squares on the spots the indexer took. Read from the spots rather than from the
    // indexer's own matrix so that whatever per-frame refinement it applied is already in it, and so
    // that no convention has to be shared with it.
    struct Orientation {
        std::array<Coord, 3> row;   // hkl = {row[0]*q, row[1]*q, row[2]*q}
    };

    bool FitOrientation(const std::vector<SpotToSave> &spots, const DiffractionGeometry &geometry,
                        Orientation &fit) {
        double n[3][3] = {};                 // sum q q^T
        double rhs[3][3] = {};               // sum q h, one column per Miller index
        size_t used = 0;
        for (const auto &s : spots) {
            // The first lattice only: on a multi-lattice frame the other lattices' Miller indices
            // belong to a different orientation and fitting one map through both gives neither.
            if (!s.indexed || s.lattice > 0)
                continue;
            const Coord q = s.ReciprocalCoord(geometry);
            const double h[3] = {static_cast<double>(s.h), static_cast<double>(s.k),
                                 static_cast<double>(s.l)};
            for (int i = 0; i < 3; i++) {
                for (int j = 0; j < 3; j++)
                    n[i][j] += static_cast<double>(q[i]) * q[j];
                for (int j = 0; j < 3; j++)
                    rhs[i][j] += static_cast<double>(q[i]) * h[j];
            }
            used++;
        }
        if (used < HARMONIC_MIN_INDEXED_SPOTS)
            return false;

        // Invert the 3x3 normal matrix by cofactors. It is singular when the frame's spots span fewer
        // than three directions in reciprocal space, which no real frame does but a degenerate one can.
        double c[3][3];
        c[0][0] = n[1][1] * n[2][2] - n[1][2] * n[2][1];
        c[0][1] = n[0][2] * n[2][1] - n[0][1] * n[2][2];
        c[0][2] = n[0][1] * n[1][2] - n[0][2] * n[1][1];
        c[1][0] = n[1][2] * n[2][0] - n[1][0] * n[2][2];
        c[1][1] = n[0][0] * n[2][2] - n[0][2] * n[2][0];
        c[1][2] = n[0][2] * n[1][0] - n[0][0] * n[1][2];
        c[2][0] = n[1][0] * n[2][1] - n[1][1] * n[2][0];
        c[2][1] = n[0][1] * n[2][0] - n[0][0] * n[2][1];
        c[2][2] = n[0][0] * n[1][1] - n[0][1] * n[1][0];
        const double det = n[0][0] * c[0][0] + n[0][1] * c[1][0] + n[0][2] * c[2][0];
        const double scale = n[0][0] + n[1][1] + n[2][2];
        if (!(std::fabs(det) > 1e-12 * scale * scale * scale))
            return false;

        for (int j = 0; j < 3; j++) {
            float v[3];
            for (int i = 0; i < 3; i++) {
                double sum = 0;
                for (int k = 0; k < 3; k++)
                    sum += c[i][k] * rhs[k][j];
                v[i] = static_cast<float>(sum / det);
            }
            fit.row[j] = Coord(v[0], v[1], v[2]);
        }
        return true;
    }

    // How far the scaled spot is from the nearest lattice point, as the largest of the three Miller
    // indices' distances to an integer.
    float IndexResidual(const Orientation &fit, const Coord &q, float scale) {
        float worst = 0;
        for (int j = 0; j < 3; j++) {
            const float f = scale * (fit.row[j] * q);
            worst = std::max(worst, std::fabs(f - std::round(f)));
        }
        return worst;
    }
}

void HarmonicEvidence::Add(const HarmonicEvidence &other) {
    frames += other.frames;
    tested += other.tested;
    hits += other.hits;
    null_hits += other.null_hits;
    hit_d_sum_A += other.hit_d_sum_A;
    indexed_d_sum_A += other.indexed_d_sum_A;
    indexed += other.indexed;
}

void AddHarmonicEvidence(const std::vector<SpotToSave> &spots, const DiffractionGeometry &geometry,
                         HarmonicEvidence &evidence) {
    Orientation fit;
    if (!FitOrientation(spots, geometry, fit))
        return;

    evidence.frames++;
    for (const auto &s : spots) {
        const Coord q = s.ReciprocalCoord(geometry);
        const float d_A = 1.0f / std::max(q.Length(), 1e-6f);
        if (s.indexed) {
            evidence.indexed++;
            evidence.indexed_d_sum_A += d_A;
            continue;
        }
        if (IndexResidual(fit, q, 1.0f) < HARMONIC_NEAR_MISS_TOLERANCE)
            continue;
        evidence.tested++;
        if (IndexResidual(fit, q, HARMONIC_SCALE) < HARMONIC_INDEX_TOLERANCE) {
            evidence.hits++;
            evidence.hit_d_sum_A += d_A;
        }
        if (IndexResidual(fit, q, HARMONIC_SCALE_NULL) < HARMONIC_INDEX_TOLERANCE)
            evidence.null_hits++;
    }
}

HarmonicContaminationResult HarmonicFromEvidence(const HarmonicEvidence &evidence) {
    HarmonicContaminationResult result;
    result.frames = evidence.frames;
    result.tested = evidence.tested;
    result.hits = evidence.hits;
    // Frames enough for the question to have been asked of this crystal at all. A run that indexed
    // this much and found nothing left to test has answered it: there is no second population.
    if (evidence.frames < 20 || evidence.indexed < 1000)
        return result;

    result.measurable = true;
    if (evidence.tested > 0) {
        result.fraction = 100.0f * static_cast<float>(evidence.hits)
                          / static_cast<float>(evidence.tested);
        result.null_fraction = 100.0f * static_cast<float>(evidence.null_hits)
                               / static_cast<float>(evidence.tested);
    }
    if (evidence.hits > 0)
        result.mean_d_A = static_cast<float>(evidence.hit_d_sum_A / static_cast<double>(evidence.hits));
    if (evidence.indexed > 0)
        result.mean_indexed_d_A = static_cast<float>(evidence.indexed_d_sum_A
                                                     / static_cast<double>(evidence.indexed));
    // Ten per cent of the tested spots, five times whatever the same test finds at a scale that is
    // not a harmonic, and enough spots for the rate to be a rate. Measured over datasets with no
    // reported contamination the harmonic rate runs at or below the null's own per cent or so, and
    // on a set whose depositors reported a harmonic it is ninety.
    result.detected = evidence.hits >= 200
                      && result.fraction >= 10.0f
                      && result.fraction >= 5.0f * std::max(result.null_fraction, 0.5f);
    return result;
}