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

#include <algorithm>
#include <cmath>
#include <string>
#include <tuple>
#include <vector>

#include "Rugnux.h"
#include "RugnuxPipeline.h"
#include "../common/JFJochMath.h"   // PI - M_PI is not portable to MSVC
#include "../image_analysis/indexing/AnalyzeIndexing.h"   // SPOT_BUDGET_SIGNIFICANCE_Z

// The far-end check. LOG ONLY: it decides nothing and changes nothing the run writes.
//
// The decisions before pass 2 - the lattice, the cell, the rotation scale, the orientation the
// sweep is predicted from, the shadow mask - are taken on the first part of the sweep
// (PrepassEnd), so that pass 1 can run during collection. Pass 2 then integrates the whole sweep,
// and what it measured on the rest is evidence those decisions never saw. Here the sweep is cut
// into blocks of FAR_END_BLOCK_DEG, a handful of per-frame quantities pass 2 already has in memory
// are averaged per block, and each quantity's level over the rest of the sweep is compared with its
// level on the first part, the block means of each part being its samples (Welch's t): what a part's
// blocks scatter by - orientation, dose, counting noise alike - is that part's noise, so one odd
// block is not a shift. A quantity is out of line when its two-sided p, times the number of
// quantities compared, is below the significance the run tests everything else at
// (SPOT_BUDGET_SIGNIFICANCE_Z, 0.1 %): an unchanged sweep then says "hold" in all but one run in a
// thousand. No number in it is set by hand beyond the block width.
//
// The quantities, and the decision each one speaks to:
//  - the share of frames that index on their own, and the share of each frame's spots on the
//    lattice: the lattice (a cell that grows with the dose, a crystal that slips);
//  - the cell scale the strong reflections ask for, d at the observed centroid over d predicted:
//    the cell;
//  - the residual rotation between each frame's refined orientation and the sweep's lattice turned
//    to that frame, about the spindle and off it: the rotation scale and the orientation;
//  - the background: the shadow mask (hardware that shadows more, or less, as the sample turns).

namespace {
    constexpr double FAR_END_BLOCK_DEG = 10.0;

    // P(|T| > t) for Student's t on an integer number of degrees of freedom, in closed form
    // (Abramowitz and Stegun 26.7.3 and 26.7.4).
    double TwoSidedP(double t, int dof) {
        if (!std::isfinite(t))
            return std::isnan(t) ? 1.0 : 0.0;
        const double theta = std::atan(std::abs(t) / std::sqrt(static_cast<double>(dof)));
        const double c2 = std::cos(theta) * std::cos(theta);
        double sum = 1.0, term = 1.0;
        if (dof % 2 == 1) {
            for (int j = 1; 2 * j + 1 <= dof - 2; j++)
                sum += term *= c2 * (2.0 * j) / (2.0 * j + 1.0);
            const double a = dof == 1 ? theta : theta + std::sin(theta) * std::cos(theta) * sum;
            return std::max(0.0, 1.0 - 2.0 * a / PI);
        }
        for (int j = 1; 2 * j <= dof - 2; j++)
            sum += term *= c2 * (2.0 * j - 1.0) / (2.0 * j);
        return std::max(0.0, 1.0 - std::sin(theta) * sum);
    }

    struct Block {
        double mean = NAN;
        double se = NAN;
        int frames = 0;
    };

    Block Summarise(const std::vector<double> &v, int first, int last) {
        Block b;
        double sum = 0.0, sum2 = 0.0;
        for (int i = first; i < last; i++)
            if (std::isfinite(v[i])) {
                sum += v[i];
                sum2 += v[i] * v[i];
                b.frames++;
            }
        if (b.frames < 2)
            return b;
        b.mean = sum / b.frames;
        const double var = std::max(0.0, (sum2 - b.frames * b.mean * b.mean) / (b.frames - 1));
        b.se = std::sqrt(var / b.frames);
        return b;
    }
}

void Rugnux::FarEndCheck(PipelineLocals &p, const RotationIndexerResult &rot) {
    auto &logger = p.logger;
    const int n = p.images_to_process;
    const int early = p.prepass_images;
    if (!rot.axis || !p.indexer)
        return;
    if (early >= n) {
        logger.Info("far-end check: the decisions read the whole sweep");
        return;
    }
    const GoniometerAxis &gon = *rot.axis;
    const double step = std::abs(gon.GetAngle_deg(1.0f) - gon.GetAngle_deg(0.0f));
    if (!(step > 0.0))
        return;
    const int block = std::max(1, static_cast<int>(std::lround(FAR_END_BLOCK_DEG / step)));

    std::vector<float> indexed, spots, spots_indexed, bkg;
    p.plots.GetPlotRaw(indexed, PlotType::IndexingRate, "");
    p.plots.GetPlotRaw(spots, PlotType::SpotCount, "");
    p.plots.GetPlotRaw(spots_indexed, PlotType::SpotCountIndexed, "");
    p.plots.GetPlotRaw(bkg, PlotType::BkgEstimate, "");
    const auto at = [](const std::vector<float> &v, int i) {
        return i < static_cast<int>(v.size()) ? static_cast<double>(v[i]) : NAN;
    };

    enum { INDEXED, ON_LATTICE, CELL, ABOUT_SPINDLE, OFF_SPINDLE, BACKGROUND, N_STATS };
    struct Stat { const char *what; const char *decision; double unit; const char *suffix; const char *tag; };
    const Stat stats[N_STATS] = {
        {"frames indexing on their own", "the lattice", 100.0, "%", "indexed"},
        {"spots on the lattice", "the lattice", 100.0, "%", "on-lattice"},
        {"cell scale of the strong reflections (d observed / d predicted - 1)", "the cell", 100.0, "%", "cell"},
        {"orientation residual about the spindle", "the rotation scale", 1.0, " deg", "about-spindle"},
        {"orientation residual off the spindle", "the orientation", 1.0, " deg", "off-spindle"},
        {"background", "the shadow mask", 1.0, "", "background"},
    };
    std::vector<std::vector<double>> frame(N_STATS, std::vector<double>(n, NAN));
    const Coord spindle = gon.GetAxis().Normalize();
    const auto &outcomes = p.indexer->GetIntegrationOutcome();
    for (int o = 0; o < n; o++) {
        frame[INDEXED][o] = at(indexed, o);
        if (at(spots, o) > 0.0)
            frame[ON_LATTICE][o] = at(spots_indexed, o) / at(spots, o);
        frame[BACKGROUND][o] = at(bkg, o);
        if (o >= static_cast<int>(outcomes.size()) || outcomes[o].reflections.empty())
            continue;
        const IntegrationOutcome &out = outcomes[o];

        // The rotation R taking the sweep's lattice, turned to this frame as the per-image path
        // turns it, onto the frame's refined lattice: R = A P^-1, with P^-1's rows the reciprocal
        // vectors of P. Its antisymmetric part is the small residual rotation.
        const float angle = gon.GetAngle_deg(static_cast<float>(o)) + gon.GetWedge_deg() / 2.0f;
        const CrystalLattice pred = rot.lattice.Multiply(gon.GetTransformationAngle(-angle));
        const Coord pv[3] = {pred.Vec0(), pred.Vec1(), pred.Vec2()};
        Coord av[3] = {out.latt.Vec0(), out.latt.Vec1(), out.latt.Vec2()};
        for (int i = 0; i < 3; i++)
            if (av[i] * pv[i] < 0.0f)   // a basis sign convention, not a rotation
                av[i] = -av[i];
        const float det = pv[0] * (pv[1] % pv[2]);
        if (std::abs(det) > 0.0f) {
            const Coord rows[3] = {(pv[1] % pv[2]) / det, (pv[2] % pv[0]) / det, (pv[0] % pv[1]) / det};
            double R[3][3];
            for (int j = 0; j < 3; j++)
                for (int k = 0; k < 3; k++)
                    R[j][k] = av[0][j] * rows[0][k] + av[1][j] * rows[1][k] + av[2][j] * rows[2][k];
            const Coord w(static_cast<float>(0.5 * (R[2][1] - R[1][2])),
                          static_cast<float>(0.5 * (R[0][2] - R[2][0])),
                          static_cast<float>(0.5 * (R[1][0] - R[0][1])));
            const double about = w * spindle;
            frame[ABOUT_SPINDLE][o] = about * 180.0 / PI;
            frame[OFF_SPINDLE][o] = (w - spindle * static_cast<float>(about)).Length() * 180.0 / PI;
        }

        // The cell the strong reflections ask for: d at the observed centroid against d predicted.
        double sum = 0.0;
        int cnt = 0;
        for (const auto &r : out.reflections) {
            if (!r.observed || r.on_ice_ring || !(r.sigma > 0.0f) || r.I < 3.0f * r.sigma || !(r.d > 0.0f))
                continue;
            const float q = out.geom.DetectorToRecip(r.observed_x, r.observed_y).Length();
            if (!(q > 0.0f))
                continue;
            sum += 1.0 / (q * r.d) - 1.0;
            cnt++;
        }
        if (cnt > 0)
            frame[CELL][o] = sum / cnt;
    }

    // The blocks wholly inside the first part, and those wholly past it; a block that straddles
    // the boundary is neither, and a last block shorter than half a block is left out.
    struct Range { int first, last; };
    std::vector<Range> early_blocks, late_blocks;
    for (int first = 0; first < n; first += block) {
        const int last = std::min(n, first + block);
        if (last - first < (block + 1) / 2)
            continue;
        if (last <= early)
            early_blocks.push_back({first, last});
        else if (first >= early)
            late_blocks.push_back({first, last});
    }
    const auto deg = [&](int ordinal) { return gon.GetAngle_deg(static_cast<float>(ordinal)); };

    // Each quantity's level over the rest of the sweep against its level on the first part: the
    // block means of each part are the samples, so what one part's blocks scatter by - orientation,
    // dose, counting noise alike - is that part's noise, and a single block out of line is not a
    // shift of the part (Welch's t, its degrees of freedom rounded down).
    struct Comparison { int stat; double late, early, noise, t, p; };
    std::vector<Comparison> comparisons;
    for (int s = 0; s < N_STATS; s++) {
        const auto sample = [&](const std::vector<Range> &ranges) {
            std::vector<double> means;
            for (const auto &r : ranges)
                if (const Block b = Summarise(frame[s], r.first, r.last); std::isfinite(b.mean))
                    means.push_back(b.mean);
            double m = 0.0, v = 0.0;
            for (const double x : means)
                m += x;
            m /= std::max<size_t>(1, means.size());
            for (const double x : means)
                v += (x - m) * (x - m);
            v /= std::max<size_t>(2, means.size()) - 1;
            return std::tuple<double, double, int>{m, v, static_cast<int>(means.size())};
        };
        const auto [me, ve, ne] = sample(early_blocks);
        const auto [ml, vl, nl] = sample(late_blocks);
        if (ne < 3 || nl < 3)
            continue;
        const double ae = ve / ne, al = vl / nl;
        const double noise = std::sqrt(ae + al);
        const double t = noise > 0.0 ? (ml - me) / noise : (ml == me ? 0.0 : INFINITY);
        const double dof = noise > 0.0 ? (ae + al) * (ae + al) / (ae * ae / (ne - 1) + al * al / (nl - 1))
                                       : ne + nl - 2;
        comparisons.push_back({s, ml, me, noise, t, TwoSidedP(t, std::max(1, static_cast<int>(dof)))});
    }
    const double alpha = std::erfc(SPOT_BUDGET_SIGNIFICANCE_Z / std::sqrt(2.0));
    const double family = static_cast<double>(comparisons.size());
    const auto describe = [&](const Comparison &c) {
        const Stat &st = stats[c.stat];
        return fmt::format("{} is {:.3g}{} over the rest of the sweep against {:.3g}{} on the first part "
                           "(+- {:.2g}, t = {:.1f}, p = {:.2g} in a family of {})",
                           st.what, c.late * st.unit, st.suffix, c.early * st.unit, st.suffix,
                           c.noise * st.unit, c.t, c.p, comparisons.size());
    };
    std::vector<std::string> revisits;
    std::string by_quantity;   // each quantity's p times the family, for the record
    const Comparison *closest = nullptr;
    for (const auto &c : comparisons) {
        by_quantity += fmt::format("{}{} {:.2g}", by_quantity.empty() ? "" : ", ", stats[c.stat].tag,
                                   std::min(1.0, c.p * family));
        if (c.p * family < alpha)
            revisits.push_back(fmt::format("{} would be revisited because {}", stats[c.stat].decision,
                                           describe(c)));
        if (!closest || c.p < closest->p)
            closest = &c;
    }

    const std::string blocks = fmt::format("{} early and {} late blocks of {:.0f} deg", early_blocks.size(),
                                           late_blocks.size(), block * step);
    if (!closest)
        logger.Info("far-end check: nothing to compare ({}; the first part ends at {:.0f} deg)", blocks,
                    deg(early));
    else if (revisits.empty())
        logger.Info("far-end check: early decisions hold ({}; closest: {}; p x family by quantity: {})",
                    blocks, describe(*closest), by_quantity);
    else {
        std::string line;
        for (const auto &r : revisits)
            line += (line.empty() ? "" : "; ") + r;
        logger.Info("far-end check: {} ({}; p x family by quantity: {})", line, blocks, by_quantity);
    }
}