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

#include <algorithm>
#include <cmath>
#include <utility>

#include "../../common/ConnectedComponents.h"
#include "../../common/JFJochMath.h"
#include "AnalyzeGridScan.h"

namespace {
    constexpr float NO_VALUE = -1.0f;
}

GridScanResult AnalyzeGridScan(const ScanResult &scan,
                               const GridScanSettings &grid,
                               float beam_size_x_um,
                               float beam_size_y_um,
                               const GridScanAnalysisSettings &settings) {
    const int64_t nx = grid.GetGridSizeX_step();
    const int64_t ny = grid.GetGridSizeY_step();
    const float step_x = fabsf(grid.GetGridStepX_um());
    const float step_y = fabsf(grid.GetGridStepY_um());

    // Scatter the per-image quantities onto the display grid. Rearrange knows about snake order,
    // the vertical flag and the step signs, so nothing here has to.
    std::vector<float> protein(nx * ny, NO_VALUE);
    std::vector<float> ice(nx * ny, NO_VALUE);
    std::vector<float> res(nx * ny, NO_VALUE);
    std::vector<int64_t> image_no(nx * ny, -1);

    for (const auto &elem: scan.images) {
        if (elem.number < 0 || elem.number >= grid.GetNElem())
            continue;
        const int64_t i = grid.Rearrange(elem.number);
        protein[i] = elem.protein_score.value_or(NO_VALUE);
        ice[i] = elem.ice_score.value_or(NO_VALUE);
        res[i] = elem.res.value_or(NO_VALUE);
        image_no[i] = elem.number;
    }

    // Hysteresis. The patches are labelled at the GROW level and then only those holding at least
    // one cell at the SEED level are kept, which is the same thing as growing out of the seeds and
    // is one pass of the labeller rather than two. A cell between the two levels therefore joins a
    // crystal that already exists but can never start one, so however low the grow level is set, a
    // background that never reaches the seed level produces nothing.
    std::vector<uint8_t> grown(nx * ny);
    for (int64_t i = 0; i < nx * ny; i++)
        grown[i] = (protein[i] > settings.GetGrowScoreThreshold()) ? 1 : 0;

    // Labelled with no size cut of its own: whether a patch is big enough is no longer a plain
    // floor - a small patch survives on the strength of its diffraction - and that test needs the
    // patch's scores, which the labeller does not have.
    const std::vector<int32_t> label = LabelConnectedComponents(grown, nx, ny, 1);
    const int32_t n_label = label.empty() ? 0 : *std::max_element(label.begin(), label.end());

    GridScanResult result;
    // An unstated beam falls back to the grid step. A raster is stepped at about the beam size, so
    // the step is the best proxy there is, and it is a far better one than zero: zero tells a
    // consumer the extents are exact when they still contain a whole beam.
    result.beam_size_x_um = beam_size_x_um > 0.0f ? beam_size_x_um : step_x;
    result.beam_size_y_um = beam_size_y_um > 0.0f ? beam_size_y_um : step_y;

    for (int32_t l = 1; l <= n_label; l++) {
        std::vector<int64_t> cell;
        for (int64_t i = 0; i < nx * ny; i++) {
            if (label[i] == l)
                cell.push_back(i);
        }
        if (cell.empty())
            continue;

        // Hysteresis, the second half: a patch that never reaches the seed level is not a patch,
        // whatever its shape, so nothing about it is worth measuring.
        bool has_seed = false;
        for (int64_t i: cell)
            has_seed = has_seed || protein[i] > settings.GetProteinScoreThreshold();
        if (!has_seed)
            continue;

        const auto n = static_cast<float>(cell.size());

        // Cell centres in MICROMETRES, taken once. Every geometric quantity below is computed from
        // these and never from cell units: step_x and step_y genuinely differ (20 x 16 um is an
        // ordinary raster), so a second moment taken in cells gives the wrong axis angle.
        // The i / nx and i % nx are the row and column of a row-major index - an index split, and
        // truncation is the whole point of it; there is no precision to lose there.
        std::vector<float> px(cell.size()), py(cell.size());
        for (size_t k = 0; k < cell.size(); k++) {
            px[k] = static_cast<float>(cell[k] % nx) * step_x;
            py[k] = static_cast<float>(cell[k] / nx) * step_y;
        }

        // Centre, pulled towards the cells that diffract best. The pull uses the RANK of the
        // resolution within the blob, never its value: the best cell gets p = 1, the worst p = 0,
        // and a cell with no resolution at all gets p = 0 rather than being dropped. A salt grain
        // reporting an absurd 0.8 A is then weighted exactly like a genuine best cell, so no
        // artefact can drag the centre however extreme its number is.
        std::vector<int64_t> by_res;
        for (int64_t i: cell) {
            if (res[i] > 0)
                by_res.push_back(i);
        }
        std::sort(by_res.begin(), by_res.end(), [&](int64_t a, int64_t b) { return res[a] < res[b]; });

        std::vector<float> weight(cell.size(), 1.0f);
        for (size_t r = 0; r < by_res.size(); r++) {
            const float p = (by_res.size() == 1) ? 1.0f
                    : 1.0f - static_cast<float>(r) / static_cast<float>(by_res.size() - 1);
            // cell is built in ascending grid order, so it can be searched directly
            weight[std::lower_bound(cell.begin(), cell.end(), by_res[r]) - cell.begin()] = 1.0f + 0.5f * p;
        }

        float sum_w = 0, sum_wx = 0, sum_wy = 0;
        for (size_t k = 0; k < cell.size(); k++) {
            sum_w += weight[k];
            sum_wx += weight[k] * px[k];
            sum_wy += weight[k] * py[k];
        }
        const float cx_um = sum_wx / sum_w;
        const float cy_um = sum_wy / sum_w;

        // A weighted centroid of a banana- or L-shaped blob can land outside the blob, where no
        // image was ever collected. The image number has to name a cell that exists, so snap.
        size_t nearest = 0;
        float nearest_d2 = INFINITY;
        for (size_t k = 0; k < cell.size(); k++) {
            const float d2 = (px[k] - cx_um) * (px[k] - cx_um) + (py[k] - cy_um) * (py[k] - cy_um);
            if (d2 < nearest_d2) {
                nearest_d2 = d2;
                nearest = k;
            }
        }

        float sxx = 0, syy = 0, sxy = 0;
        for (size_t k = 0; k < cell.size(); k++) {
            const float dx = px[k] - cx_um;
            const float dy = py[k] - cy_um;
            sxx += dx * dx;
            syy += dy * dy;
            sxy += dx * dy;
        }
        const float angle = 0.5f * atan2f(2 * sxy, sxx - syy);
        const float cos_a = cosf(angle);
        const float sin_a = sinf(angle);

        // Direction from the eigenvector, LENGTH from the projected extent. "How far do I scan"
        // is an extent question, and the constant taking a second moment to a length depends on
        // an assumed shape that a blob of a few cells does not have.
        float min_u = INFINITY, max_u = -INFINITY, min_v = INFINITY, max_v = -INFINITY;
        for (size_t k = 0; k < cell.size(); k++) {
            const float dx = px[k] - cx_um;
            const float dy = py[k] - cy_um;
            min_u = std::min(min_u, dx * cos_a + dy * sin_a);
            max_u = std::max(max_u, dx * cos_a + dy * sin_a);
            min_v = std::min(min_v, -dx * sin_a + dy * cos_a);
            max_v = std::max(max_v, -dx * sin_a + dy * cos_a);
        }
        // The span runs between cell centres, so one cell has to be added back. A cell is a
        // step_x by step_y rectangle, and its own width along a direction is that rectangle's
        // support width - which is the plain step only when the axis lies along the grid.
        const float cell_along_u = fabsf(step_x * cos_a) + fabsf(step_y * sin_a);
        const float cell_along_v = fabsf(step_x * sin_a) + fabsf(step_y * cos_a);

        float sum_protein = 0, sum_ice = 0, peak_protein = 0;
        for (int64_t i: cell) {
            sum_protein += protein[i];
            sum_ice += std::max(ice[i], 0.0f);
            peak_protein = std::max(peak_protein, protein[i]);
        }

        // A patch is a crystal when it is big enough to be a shape rather than a coincidence, OR
        // when it is smaller than that but the diffraction in it is decisive on its own. One cell
        // is enough where that cell is clearly protein; a weak patch still has to be a shape.
        //
        // Both halves are read over the GROWN patch: the count includes the cells hysteresis added
        // and the peak is the patch's best cell wherever it lies. That cannot let growth rescue a
        // patch no seed would have admitted, because the patch has already been discarded above
        // unless it holds a seed cell - and a seed cell is by definition the strongest kind there
        // is, so the peak of a grown patch is the peak of its seeds.
        if (static_cast<int64_t>(cell.size()) < settings.GetMinBlobCells()
            && peak_protein < settings.GetDecisiveSingleCellScore())
            continue;

        GridScanCrystal crystal;
        crystal.nx = cx_um / step_x;
        crystal.ny = cy_um / step_y;
        crystal.x_um = cx_um;
        crystal.y_um = cy_um;
        crystal.image_number = image_no[cell[nearest]];
        crystal.major_um = max_u - min_u + cell_along_u;
        crystal.minor_um = max_v - min_v + cell_along_v;
        crystal.angle_deg = angle * 180.0f / static_cast<float>(PI);

        // The angle is the axis of the larger second MOMENT, the extents are MEASURED spans, and
        // for a strongly non-convex blob the two can disagree about which axis is the longer.
        // A consumer draws a major by minor frame rotated by angle_deg, so keep both facts by
        // turning the frame a quarter turn rather than by dropping one of them.
        if (crystal.major_um < crystal.minor_um) {
            std::swap(crystal.major_um, crystal.minor_um);
            crystal.angle_deg += 90.0f;
        }

        // atan2 returns (-pi,pi], so the half-angle is in (-pi/2,pi/2] and the quarter turn above
        // can carry it past 180; an axis has no sign, so fold it into [0,180).
        if (crystal.angle_deg < 0)
            crystal.angle_deg += 180.0f;
        if (crystal.angle_deg >= 180.0f)
            crystal.angle_deg -= 180.0f;

        // The MEAN protein score, not the peak: the score saturates, so the peak is 1.0 for every
        // real crystal and ranks nothing. The mean stays a detection confidence and compares.
        crystal.score = sum_protein / n;
        crystal.ice_score = sum_ice / n;
        crystal.peak_score = peak_protein;

        // The 25th percentile, not the minimum: the single best cell in a blob is precisely where
        // a salt spot or a hot pixel shows up.
        if (!by_res.empty()) {
            const auto q = static_cast<size_t>(0.25 * static_cast<double>(by_res.size() - 1) + 0.5);
            crystal.res_A = res[by_res[q]];
        }
        crystal.n_images = static_cast<int64_t>(cell.size());

        result.crystals.push_back(crystal);
    }

    std::sort(result.crystals.begin(), result.crystals.end(),
              [](const GridScanCrystal &a, const GridScanCrystal &b) { return a.score > b.score; });

    // Kept to the best few only where a caller asked for that; the sort above is what makes the
    // ones it keeps the right ones.
    // Unset means no cap: a crystal found and then dropped is information the caller cannot recover.
    if (const auto cap = settings.GetMaxCrystals();
        cap.has_value() && result.crystals.size() > static_cast<size_t>(*cap))
        result.crystals.resize(*cap);

    return result;
}