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

#include <catch2/catch_all.hpp>

#include <cmath>
#include <random>
#include <vector>

#include "../common/DiffractionGeometry.h"
#include "../image_analysis/SensorAbsorption.h"
#include "../rugnux/SpotWidth.h"

namespace {
    constexpr int W = 3000, H = 3000, C = 1500;
    constexpr double DIST_MM = 100.0, PIXEL_MM = 0.075, LAMBDA_A = 1.0;
    constexpr double THICKNESS_UM = 450.0;
    constexpr double ANGULAR_SIGMA = 1.5e-4;   // rad: divergence and crystal, the same in every direction

    // Spots on an untilted detector, each an elliptical Gaussian with its long axis along the radius:
    // the isotropic angular width, the bandwidth's radial streak 2 tan(theta) sigma, and the sensor
    // parallax - a conversion depth z moves the photon by z tan(psi) along the radius, and on an
    // untilted detector psi = 2theta. Worked out here from the flat-detector geometry, independently
    // of the estimator's own Jacobians: R = D tan(2theta), so one radian of 2theta is D/cos^2(2theta)
    // along the radius and one radian across the scattering plane D/cos(2theta) across it. Painted
    // with 4x4 sub-pixel sampling so each pixel carries the box it integrates, and Poisson noise.
    void PaintSpots(std::vector<int32_t> &image, std::vector<DiffractionSpot> &spots, double bw_fwhm,
                    int offset, std::mt19937 &rng) {
        const double sigma_bw = bw_fwhm / 2.3548;
        const double L_um = sensor_absorption::AttenuationLength_um("Si", LAMBDA_A);
        std::vector<double> mean(image.size(), 3.0);
        for (int cy = 40 + offset; cy < H - 40; cy += 60)
            for (int cx = 40 + offset; cx < W - 40; cx += 60) {
                const double rx = (cx - C) * PIXEL_MM, ry = (cy - C) * PIXEL_MM, r = std::hypot(rx, ry);
                if (r < 8.0) continue;
                const double tt = std::atan2(r, DIST_MM), c2 = std::cos(tt);
                const double jr = DIST_MM / (PIXEL_MM * c2 * c2), jt = DIST_MM / (PIXEL_MM * c2);
                const double depth_var = sensor_absorption::ConversionDepthVariance_um2(L_um * c2, THICKNESS_UM);
                const double par = depth_var * std::tan(tt) * std::tan(tt) / (PIXEL_MM * 1000.0 * PIXEL_MM * 1000.0);
                const double streak = 2.0 * std::tan(tt / 2.0) * sigma_bw;
                const double var_u = jr * jr * (ANGULAR_SIGMA * ANGULAR_SIGMA + streak * streak) + par;
                const double var_v = jt * jt * ANGULAR_SIGMA * ANGULAR_SIGMA;
                const double ux = rx / r, uy = ry / r;
                const double total = 20000.0;
                const double norm = total / (2.0 * M_PI * std::sqrt(var_u * var_v) * 16.0);
                for (int dy = -14; dy <= 14; dy++)
                    for (int dx = -14; dx <= 14; dx++) {
                        double v = 0.0;
                        for (int sy = 0; sy < 4; sy++)
                            for (int sx = 0; sx < 4; sx++) {
                                const double px = dx - 0.375 + 0.25 * sx, py = dy - 0.375 + 0.25 * sy;
                                const double u = px * ux + py * uy, t = -px * uy + py * ux;
                                v += std::exp(-0.5 * (u * u / var_u + t * t / var_v));
                            }
                        mean[static_cast<size_t>(cy + dy) * W + cx + dx] += norm * v;
                    }
                spots.emplace_back(static_cast<uint32_t>(cx), static_cast<uint32_t>(cy),
                                   static_cast<int64_t>(total));
            }
        for (size_t i = 0; i < image.size(); i++)
            image[i] = std::poisson_distribution<int32_t>(mean[i])(rng);
    }

    std::optional<spot_width::BandwidthEstimate> Estimate(double bw_fwhm) {
        DiffractionGeometry geometry;
        geometry.BeamX_pxl(C).BeamY_pxl(C).DetectorDistance_mm(DIST_MM).PixelSize_mm(PIXEL_MM)
                .Wavelength_A(LAMBDA_A);
        std::mt19937 rng(7);
        std::vector<spot_width::FluxCurve> curves;
        for (int offset : {0, 20, 40}) {
            ImagePreprocessorBuffer image(static_cast<size_t>(W) * H);
            std::vector<DiffractionSpot> spots;
            PaintSpots(image.getBuffer(), spots, bw_fwhm, offset, rng);
            MeasureSpotFluxCurves(image, W, H, geometry, spots, curves);
        }
        return spot_width::EstimateBandwidth(curves, sensor_absorption::AttenuationLength_um("Si", LAMBDA_A),
                                             THICKNESS_UM, PIXEL_MM * 1000.0);
    }
}

// A 0.45 % FWHM bandwidth - a multilayer's - on a thick silicon sensor, whose parallax elongates the
// spots along the radius as well and grows with angle much as the bandwidth does. The estimator has
// to take the parallax out from the sensor's physics and return the bandwidth, significantly.
TEST_CASE("BandwidthEstimate_RecoversBandwidthPastParallax", "[SpotWidth]") {
    const auto e = Estimate(0.0045);
    REQUIRE(e.has_value());
    CHECK(e->spots >= spot_width::BANDWIDTH_MIN_SPOTS);
    CHECK(e->fwhm == Catch::Approx(0.0045).margin(0.0006));
    CHECK(e->z > spot_width::BANDWIDTH_Z);
}

// The same spots from a monochromatic beam: the parallax alone must not read as a bandwidth.
TEST_CASE("BandwidthEstimate_MonochromaticIsNotSignificant", "[SpotWidth]") {
    const auto e = Estimate(0.0);
    REQUIRE(e.has_value());
    CHECK(e->z < spot_width::BANDWIDTH_Z);
    CHECK(std::abs(e->fwhm) < 0.0015);
}