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

#include "ShadowFinder.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <queue>
#include <type_traits>

#include "../../common/JFJochException.h"

// ---------------------------------------------------------------------------------
// Small binary-image helpers on a width*height frame stored row-major as char (0/1).
// All run once, at GetMask() time. The BFS forms keep them O(pixels) rather than
// O(pixels * radius), so a radius-14 dilation is still a single sweep.
// ---------------------------------------------------------------------------------
namespace {

// 8-connected dilation by `r` pixels (Chebyshev), via a multi-source BFS.
std::vector<char> Dilate(const std::vector<char> &in, int W, int H, int r) {
    if (r <= 0)
        return in;
    std::vector<int> dist(in.size(), -1);
    std::queue<int> q;
    for (size_t i = 0; i < in.size(); i++)
        if (in[i]) { dist[i] = 0; q.push(static_cast<int>(i)); }
    while (!q.empty()) {
        const int i = q.front(); q.pop();
        if (dist[i] >= r)
            continue;
        const int y = i / W, x = i % W;
        for (int dy = -1; dy <= 1; dy++)
            for (int dx = -1; dx <= 1; dx++) {
                const int yy = y + dy, xx = x + dx;
                if (yy < 0 || yy >= H || xx < 0 || xx >= W)
                    continue;
                const int j = yy * W + xx;
                if (dist[j] < 0) { dist[j] = dist[i] + 1; q.push(j); }
            }
    }
    std::vector<char> out(in.size());
    for (size_t i = 0; i < out.size(); i++)
        out[i] = (dist[i] >= 0) ? 1 : 0;
    return out;
}

// Erosion by `r` = dilation of the complement (image border counts as outside).
std::vector<char> Erode(const std::vector<char> &in, int W, int H, int r) {
    std::vector<char> comp(in.size());
    for (size_t i = 0; i < in.size(); i++)
        comp[i] = !in[i];
    const auto grown = Dilate(comp, W, H, r);
    std::vector<char> out(in.size());
    for (size_t i = 0; i < out.size(); i++)
        out[i] = !grown[i];
    return out;
}

// Pixels of `passable` reachable from any of `seeds` (8-connected flood).
std::vector<char> Flood(const std::vector<char> &passable, int W, int H, const std::vector<int> &seeds) {
    std::vector<char> visited(passable.size(), 0);
    std::queue<int> q;
    for (const int s : seeds)
        if (s >= 0 && s < static_cast<int>(passable.size()) && passable[s] && !visited[s]) {
            visited[s] = 1; q.push(s);
        }
    while (!q.empty()) {
        const int i = q.front(); q.pop();
        const int y = i / W, x = i % W;
        for (int dy = -1; dy <= 1; dy++)
            for (int dx = -1; dx <= 1; dx++) {
                const int yy = y + dy, xx = x + dx;
                if (yy < 0 || yy >= H || xx < 0 || xx >= W)
                    continue;
                const int j = yy * W + xx;
                if (passable[j] && !visited[j]) { visited[j] = 1; q.push(j); }
            }
    }
    return visited;
}

// Fill holes: background not reachable from the image border becomes region.
std::vector<char> FillHoles(const std::vector<char> &region, int W, int H) {
    std::vector<char> bg_visited(region.size(), 0);
    std::queue<int> q;
    auto push = [&](int i) { if (!region[i] && !bg_visited[i]) { bg_visited[i] = 1; q.push(i); } };
    for (int x = 0; x < W; x++) { push(x); push((H - 1) * W + x); }
    for (int y = 0; y < H; y++) { push(y * W); push(y * W + W - 1); }
    while (!q.empty()) {
        const int i = q.front(); q.pop();
        const int y = i / W, x = i % W;
        for (int dy = -1; dy <= 1; dy++)
            for (int dx = -1; dx <= 1; dx++) {
                const int yy = y + dy, xx = x + dx;
                if (yy < 0 || yy >= H || xx < 0 || xx >= W)
                    continue;
                const int j = yy * W + xx;
                if (!region[j] && !bg_visited[j]) { bg_visited[j] = 1; q.push(j); }
            }
    }
    std::vector<char> out = region;
    for (size_t i = 0; i < out.size(); i++)
        if (!region[i] && !bg_visited[i])
            out[i] = 1;
    return out;
}

// Median of `values` per integer radius, over the pixels flagged in `use`.
std::vector<float> RingMedian(const std::vector<float> &values, const std::vector<char> &use,
                              const std::vector<int> &radius, int max_radius) {
    std::vector<std::vector<float>> bins(max_radius + 1);
    for (size_t i = 0; i < values.size(); i++)
        if (use[i])
            bins[radius[i]].push_back(values[i]);
    std::vector<float> median(max_radius + 1, 0.0f);
    for (int r = 0; r <= max_radius; r++) {
        auto &b = bins[r];
        if (!b.empty()) {
            const size_t k = b.size() / 2;
            std::nth_element(b.begin(), b.begin() + k, b.end());
            median[r] = b[k];
        }
    }
    return median;
}

// Fraction of each integer-radius ring that is flagged in `blocked`.
std::vector<float> RingFraction(const std::vector<char> &blocked, const std::vector<int> &radius, int max_radius) {
    std::vector<int64_t> num(max_radius + 1, 0), den(max_radius + 1, 0);
    for (size_t i = 0; i < blocked.size(); i++) {
        den[radius[i]]++;
        if (blocked[i]) num[radius[i]]++;
    }
    std::vector<float> frac(max_radius + 1, 0.0f);
    for (int r = 0; r <= max_radius; r++)
        frac[r] = den[r] ? static_cast<float>(num[r]) / static_cast<float>(den[r]) : 0.0f;
    return frac;
}

}  // namespace

// ---------------------------------------------------------------------------------

ShadowFinder::ShadowFinder(const DiffractionExperiment &experiment, ShadowFinderSettings in_settings)
    : width(static_cast<int>(experiment.GetXPixelsNumConv())),
      height(static_cast<int>(experiment.GetYPixelsNumConv())),
      beam_x(experiment.GetBeamX_pxl()),
      beam_y(experiment.GetBeamY_pxl()),
      settings(in_settings),
      max_value(static_cast<size_t>(width) * height, 0),
      sum_value(static_cast<size_t>(width) * height, 0),
      valid_count(static_cast<size_t>(width) * height, 0) {}

template<class T>
void ShadowFinder::Add(const T *ptr) {
    // The pixel type's sentinel extreme marks "no data" (module gap / masked): the
    // preprocessor/writer stores INT*_MIN for signed and UINT*_MAX for unsigned. For
    // signed types the opposite extreme (INT*_MAX) is a genuine saturated value and is
    // kept, so a saturated reflection still registers as bright.
    T masked;
    if constexpr (std::is_signed_v<T>)
        masked = std::numeric_limits<T>::min();
    else
        masked = std::numeric_limits<T>::max();

    std::unique_lock ul(m);
    for (size_t i = 0; i < max_value.size(); i++) {
        const T v = ptr[i];
        if (v == masked)
            continue;
        const int32_t vi = static_cast<int32_t>(v);
        if (valid_count[i] == 0 || vi > max_value[i])
            max_value[i] = vi;
        sum_value[i] += vi;
        valid_count[i]++;
    }
    frames++;
}

void ShadowFinder::AddImage(const DataMessage &data, std::vector<uint8_t> buffer) {
    if (static_cast<size_t>(data.image.GetWidth()) * data.image.GetHeight() != max_value.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "ShadowFinder: image size does not match the detector");

    const auto ptr = data.image.GetUncompressedPtr(buffer);
    switch (data.image.GetMode()) {
        case CompressedImageMode::Int8:   Add(reinterpret_cast<const int8_t   *>(ptr)); break;
        case CompressedImageMode::Uint8:  Add(reinterpret_cast<const uint8_t  *>(ptr)); break;
        case CompressedImageMode::Int16:  Add(reinterpret_cast<const int16_t  *>(ptr)); break;
        case CompressedImageMode::Uint16: Add(reinterpret_cast<const uint16_t *>(ptr)); break;
        case CompressedImageMode::Int32:  Add(reinterpret_cast<const int32_t  *>(ptr)); break;
        case CompressedImageMode::Uint32: Add(reinterpret_cast<const uint32_t *>(ptr)); break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "ShadowFinder: unsupported image mode");
    }
}

uint32_t ShadowFinder::GetFrameCount() const {
    std::unique_lock ul(m);
    return frames;
}

std::vector<uint32_t> ShadowFinder::GetMask() const {
    std::unique_lock ul(m);

    const int W = width, H = height;
    const int N = W * H;
    const ShadowFinderSettings &S = settings;

    std::vector<uint32_t> mask(N, 0);
    if (frames == 0)
        return mask;

    // --- mean projection, per-pixel validity and radius from the beam centre ---
    std::vector<float> mean(N, 0.0f);
    std::vector<char>  valid(N, 0);
    std::vector<int>   radius(N, 0);
    int max_radius = 0;
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            const int i = y * W + x;
            if (valid_count[i] > 0) {
                mean[i] = static_cast<float>(static_cast<double>(sum_value[i]) / valid_count[i]);
                valid[i] = 1;
            }
            const double dx = x - beam_x, dy = y - beam_y;
            const int r = static_cast<int>(std::lround(std::sqrt(dx * dx + dy * dy)));
            radius[i] = r;
            if (r > max_radius) max_radius = r;
        }

    // --- robust radial baseline; iterate to keep the shadow out of its own baseline ---
    std::vector<float> ratio(N, 1.0f);
    std::vector<char>  excluded(N, 0);
    for (int iter = 0; iter < 3; iter++) {
        std::vector<char> use(N);
        for (int i = 0; i < N; i++)
            use[i] = valid[i] && !excluded[i];
        const auto baseline = RingMedian(mean, use, radius, max_radius);
        for (int i = 0; i < N; i++)
            if (valid[i])
                ratio[i] = mean[i] / std::max(baseline[radius[i]], 1e-6f);
        for (int i = 0; i < N; i++)
            excluded[i] = valid[i] && ratio[i] < S.shadow_ratio;
    }

    // --- shadow core: low-ratio pixels connected to the beam centre (bridging gaps) ---
    std::vector<char> low(N);
    for (int i = 0; i < N; i++)
        low[i] = valid[i] && ratio[i] < S.shadow_ratio;

    const std::vector<char> grown = Dilate(low, W, H, S.bridge_px);
    std::vector<int> seeds;   // a small disk at the beam centre
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            const double dx = x - beam_x, dy = y - beam_y;
            if (dx * dx + dy * dy < 4.0 * 4.0)
                seeds.push_back(y * W + x);
        }
    const std::vector<char> connected = Flood(grown, W, H, seeds);
    std::vector<char> core(N);
    for (int i = 0; i < N; i++)
        core[i] = low[i] && connected[i];

    // --- real reflections: any pixel that recorded signal is never masked. Require a
    //     small cluster so a single-frame zinger does not count as a reflection. ---
    std::vector<char> lit(N, 0);
    for (int i = 0; i < N; i++)
        lit[i] = (valid_count[i] > 0) && (max_value[i] >= static_cast<int32_t>(S.min_reflection));
    std::vector<char> reflection(N, 0);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            const int i = y * W + x;
            if (!lit[i]) continue;
            int neighbours = 0;
            for (int dy = -1; dy <= 1; dy++)
                for (int dx = -1; dx <= 1; dx++) {
                    const int yy = y + dy, xx = x + dx;
                    if ((dx || dy) && yy >= 0 && yy < H && xx >= 0 && xx < W && lit[yy * W + xx])
                        neighbours++;
                }
            reflection[i] = (neighbours >= 2);
        }

    // --- central low-res disk: the fully-blocked region about the beam centre. Sized by
    //     the azimuthal blocked fraction (a disk blocks ~every azimuth; a thin arm or
    //     gap does not), and capped just inside the innermost reflection. ---
    std::vector<char> blocked(N);
    for (int i = 0; i < N; i++)
        blocked[i] = (valid_count[i] == 0) || low[i];
    const auto blocked_frac = RingFraction(blocked, radius, max_radius);

    int disk_radius = 0;
    {
        float head = 0.0f; int head_n = 0;
        for (int r = 0; r <= std::min(5, max_radius); r++) { head += blocked_frac[r]; head_n++; }
        if (head_n > 0 && head / head_n >= 0.65f) {   // the beam centre is behind a disk
            disk_radius = max_radius;
            for (int r = 1; r <= max_radius; r++)
                if (blocked_frac[r] < 0.55f) { disk_radius = r; break; }
        }
    }
    int reflection_radius = max_radius + 1;   // innermost reflection (ignore the very centre)
    for (int i = 0; i < N; i++)
        if (reflection[i] && radius[i] > 12 && radius[i] < reflection_radius)
            reflection_radius = radius[i];
    if (disk_radius > reflection_radius - 4)
        disk_radius = reflection_radius - 4;
    if (disk_radius < 0)
        disk_radius = 0;

    // --- assemble: core + disk, grow the soft penumbra, round, fill the disk interior ---
    std::vector<char> region(N);
    for (int i = 0; i < N; i++)
        region[i] = core[i] || (disk_radius > 0 && radius[i] < disk_radius);

    const std::vector<char> near = Dilate(region, W, H, S.penumbra_max_px);
    for (int i = 0; i < N; i++)
        if (near[i] && valid[i] && ratio[i] < S.penumbra_ratio)
            region[i] = 1;

    region = Erode(Dilate(region, W, H, 2), W, H, 2);   // close: round the boundary
    region = FillHoles(region, W, H);

    // Expose recorded reflections - done last, with no fill afterwards, so a spot the
    // geometry still covered is given back rather than re-enclosed.
    const std::vector<char> reflection_grown = Dilate(reflection, W, H, 1);
    for (int i = 0; i < N; i++)
        if (reflection_grown[i])
            region[i] = 0;

    for (int i = 0; i < N; i++)
        mask[i] = region[i] ? 1 : 0;
    return mask;
}