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

#include "ModuleDiagnostics.h"

#include <algorithm>
#include <cmath>
#include <random>

#include <Eigen/Dense>

#include "../common/JFJochMath.h"
#include "../common/Logger.h"
#include "../image_analysis/SensorAbsorption.h"

namespace {

constexpr double RAD2DEG = 180.0 / PI;

struct Fit {
    bool ok = false;
    std::vector<ModuleShift> modules;
    GeometryRefinerResult r;
};

Fit Place(const DiffractionGeometry &geom, const UnitCell &cell, const std::vector<GeomRefineFrame> &frames,
          const std::vector<ModuleShift> &start, GeometryRefinerSettings st, bool refine) {
    st.modules = start;
    st.refine_modules = refine;
    Fit f;
    f.r = RefineGlobalGeometry(geom, cell, frames, st);
    f.ok = f.r.ok;
    f.modules = f.r.modules;
    return f;
}

// Weighted straight line y = a + b x, with b's standard error and the weighted correlation.
struct Line { double a = 0, b = 0, se_b = 0, corr = 0; };

Line FitLine(const std::vector<double> &x, const std::vector<double> &y, const std::vector<double> &w) {
    double sw = 0, sx = 0, sy = 0;
    for (size_t i = 0; i < x.size(); i++) { sw += w[i]; sx += w[i] * x[i]; sy += w[i] * y[i]; }
    Line l;
    if (sw <= 0 || x.size() < 3)
        return l;
    const double mx = sx / sw, my = sy / sw;
    double sxx = 0, syy = 0, sxy = 0;
    for (size_t i = 0; i < x.size(); i++) {
        sxx += w[i] * (x[i] - mx) * (x[i] - mx);
        syy += w[i] * (y[i] - my) * (y[i] - my);
        sxy += w[i] * (x[i] - mx) * (y[i] - my);
    }
    if (sxx <= 0)
        return l;
    l.b = sxy / sxx;
    l.a = my - l.b * mx;
    l.corr = syy > 0 ? sxy / std::sqrt(sxx * syy) : 0;
    double rss = 0;
    for (size_t i = 0; i < x.size(); i++) {
        const double e = y[i] - l.a - l.b * x[i];
        rss += w[i] * e * e;
    }
    // Weights normalised to the number of points, so the error is that of an ordinary fit of n points.
    const double n = static_cast<double>(x.size());
    l.se_b = std::sqrt(rss / sw * n / (n - 2) / (sxx / sw * n));
    return l;
}

// Which modules count: those the placement moved (enough spots), excluding the unmoved ones.
std::vector<size_t> Moved(const GeometryRefinerResult &r, const GeometryRefinerSettings &st) {
    std::vector<size_t> out;
    for (size_t m = 0; m < r.modules.size(); m++)
        if (r.modules[m].spots >= st.min_module_spots)
            out.push_back(m);
    return out;
}

void LogShifts(Logger &log, const std::string &label, const std::vector<ModuleShift> &centres,
               const std::vector<ModuleShift> &modules, const std::vector<size_t> &which) {
    for (size_t m: which)
        log.Info("ModuleDiag {} module {} dx {:+.3f} dy {:+.3f} turn_deg {:+.4f}", label, m,
                 modules[m].dx - centres[m].dx, modules[m].dy - centres[m].dy,
                 (modules[m].rot_rad - centres[m].rot_rad) * RAD2DEG);
}

// The shift field's global patterns, fitted over the modules (weights: spots): translation, turn about
// the beam, radial scale, shear, and the two tilt-like quadratic patterns. What is left after them is
// what the modules hold on their own.
void Decompose(Logger &log, const std::string &label, double bx, double by, const std::vector<ModuleShift> &mods,
               const std::vector<size_t> &which, double lever) {
    const size_t n = which.size();
    if (n < 4)
        return;
    // Unknowns: tx, ty, a11, a12, a21, a22, tilt_y, tilt_x  (8); two equations per module.
    Eigen::MatrixXd A = Eigen::MatrixXd::Zero(2 * n, 8);
    Eigen::VectorXd b(2 * n), w(2 * n);
    for (size_t k = 0; k < n; k++) {
        const auto &m = mods[which[k]];
        const double x = (m.x0 - bx) / lever, y = (m.y0 - by) / lever;
        A(2 * k, 0) = 1; A(2 * k, 2) = x; A(2 * k, 3) = y; A(2 * k, 6) = x * x; A(2 * k, 7) = x * y;
        A(2 * k + 1, 1) = 1; A(2 * k + 1, 4) = x; A(2 * k + 1, 5) = y; A(2 * k + 1, 6) = x * y; A(2 * k + 1, 7) = y * y;
        b(2 * k) = m.dx;
        b(2 * k + 1) = m.dy;
        w(2 * k) = w(2 * k + 1) = std::sqrt(static_cast<double>(m.spots));
    }
    const Eigen::MatrixXd Aw = w.asDiagonal() * A;
    const Eigen::VectorXd bw = w.asDiagonal() * b;
    // Linear part alone, then with the tilt-like terms: how much the quadratic patterns explain.
    const Eigen::VectorXd lin = Aw.leftCols(6).colPivHouseholderQr().solve(bw);
    const Eigen::VectorXd full = Aw.colPivHouseholderQr().solve(bw);
    auto rms = [&](const Eigen::VectorXd &res) {
        double s = 0, sw = 0;
        for (Eigen::Index i = 0; i < res.size(); i++) { s += w(i) * w(i) * res(i) * res(i); sw += w(i) * w(i); }
        return std::sqrt(s / sw);
    };
    const double rms_raw = rms(b);
    const double rms_lin = rms(b - A.leftCols(6) * lin);
    const double rms_full = rms(b - A * full);
    log.Info("ModuleDiag {} decomposition: translation ({:+.3f}, {:+.3f}) px, turn {:+.4f} deg, radial scale "
             "{:+.3f} px at the lever, shear ({:+.3f}, {:+.3f}) px; tilt-like ({:+.3f}, {:+.3f}) px at the lever; "
             "rms raw {:.3f}, after linear {:.3f}, after linear+tilt {:.3f} px (lever {:.0f} px)",
             label, full(0), full(1), 0.5 * (full(4) - full(3)) / lever * RAD2DEG,
             0.5 * (full(2) + full(5)), 0.5 * (full(2) - full(5)), 0.5 * (full(3) + full(4)),
             full(6), full(7), rms_raw, rms_lin, rms_full, lever);
}

void RadialTable(Logger &log, const std::string &label, double bx, double by, const std::vector<ModuleShift> &mods,
                 const std::vector<size_t> &which, const std::vector<double> &mean_d,
                 const std::vector<double> &rms_before, const std::vector<double> &rms_after,
                 const std::vector<std::array<double, 3>> &sigma) {
    std::vector<double> r, rad, tan, mag, turn, w;
    for (size_t m: which) {
        const auto &s = mods[m];
        const double ox = s.x0 - bx, oy = s.y0 - by, rr = std::hypot(ox, oy);
        const double ux = ox / rr, uy = oy / rr;
        const double radial = s.dx * ux + s.dy * uy, tangential = -s.dx * uy + s.dy * ux;
        r.push_back(rr);
        rad.push_back(radial);
        tan.push_back(tangential);
        mag.push_back(std::hypot(s.dx, s.dy));
        turn.push_back(s.rot_rad * RAD2DEG);
        w.push_back(static_cast<double>(s.spots));
        const auto sg = m < sigma.size() ? sigma[m] : std::array<double, 3>{NAN, NAN, NAN};
        log.Info("ModuleDiag {} table module {} r_px {:.0f} mean_d {:.2f} spots {} dx {:+.3f} dy {:+.3f} radial {:+.3f} "
                 "tangential {:+.3f} turn_deg {:+.4f} sigma_dx {:.3f} sigma_dy {:.3f} sigma_turn_deg {:.4f} "
                 "rms_before {:.3f} rms_after {:.3f}", label, m, rr, m < mean_d.size() ? mean_d[m] : NAN, s.spots,
                 s.dx, s.dy, radial, tangential, s.rot_rad * RAD2DEG, sg[0], sg[1], sg[2] * RAD2DEG,
                 m < rms_before.size() ? rms_before[m] : NAN, m < rms_after.size() ? rms_after[m] : NAN);
    }
    for (const auto &[name, v]: std::vector<std::pair<std::string, std::vector<double> *>>{
                 {"radial", &rad}, {"tangential", &tan}, {"magnitude", &mag}, {"turn_deg", &turn}}) {
        const Line l = FitLine(r, *v, w);
        log.Info("ModuleDiag {} trend {} vs r: slope {:+.3e} +- {:.3e} per px ({:+.3f} +- {:.3f} px per 1000 px), "
                 "corr {:+.2f}", label, name, l.b, l.se_b, 1000 * l.b, 1000 * l.se_b, l.corr);
    }
}

}  // namespace

void RunModuleDiagnostics(const ModuleDiagnosticsInput &in) {
    Logger log("ModuleDiag");
    const auto &st = in.settings;
    const auto &placed = in.placed;
    const std::vector<size_t> which = Moved(placed, st);
    const double bx = placed.beam_x_px, by = placed.beam_y_px;
    double lever = 0;
    for (size_t m: which)
        lever += std::hypot(in.centres[m].x0 - bx, in.centres[m].y0 - by);
    lever /= std::max<size_t>(1, which.size());

    // Before: the same frames, modules held where given.
    const Fit held = Place(in.nominal, in.cell, in.frames, in.centres, st, false);

    // Bootstrap over frames: the placement's spread, per module.
    constexpr int BOOT = 20;
    std::mt19937 rng(12345);
    std::vector<std::array<double, 3>> sum(in.centres.size(), {0, 0, 0}), sum2(in.centres.size(), {0, 0, 0});
    int boots = 0;
    for (int b = 0; b < BOOT; b++) {
        std::vector<GeomRefineFrame> resample;
        std::uniform_int_distribution<size_t> pick(0, in.frames.size() - 1);
        for (size_t i = 0; i < in.frames.size(); i++)
            resample.push_back(in.frames[pick(rng)]);
        const Fit f = Place(in.nominal, in.cell, resample, in.start, st, true);
        if (!f.ok)
            continue;
        boots++;
        for (size_t m = 0; m < f.modules.size(); m++) {
            const double v[3] = {f.modules[m].dx, f.modules[m].dy, f.modules[m].rot_rad};
            for (int j = 0; j < 3; j++) { sum[m][j] += v[j]; sum2[m][j] += v[j] * v[j]; }
        }
    }
    std::vector<std::array<double, 3>> sigma(in.centres.size(), {NAN, NAN, NAN});
    if (boots >= 2)
        for (size_t m = 0; m < sigma.size(); m++)
            for (int j = 0; j < 3; j++) {
                const double mean = sum[m][j] / boots;
                sigma[m][j] = std::sqrt(std::max(0.0, sum2[m][j] / boots - mean * mean) * boots / (boots - 1));
            }
    log.Info("ModuleDiag bootstrap: {} of {} resamples placed", boots, BOOT);

    log.Info("ModuleDiag placed: beam ({:.3f}, {:.3f}) distance {:.4f} rot1 {:.5f} rot2 {:.5f} cell {:.3f} {:.3f} {:.3f}",
             placed.beam_x_px, placed.beam_y_px, placed.distance_mm, placed.rot1_rad, placed.rot2_rad,
             placed.cell.a, placed.cell.b, placed.cell.c);
    RadialTable(log, "placed", bx, by, placed.modules, which, placed.module_mean_d_A, held.r.module_capped_rms_px,
                placed.module_capped_rms_px, sigma);
    Decompose(log, "placed", bx, by, placed.modules, which, lever);

    // Low against high resolution: the placement fitted on the spots coarser than 3 A, and on the finer ones.
    for (int half = 0; half < 2; half++) {
        std::vector<GeomRefineFrame> sub = in.frames;
        for (auto &f: sub) {
            std::vector<GeomRefineSpot> keep;
            for (const auto &s: f.spots) {
                const double d = 1.0 / std::max(1e-6f, in.nominal.DetectorToRecip(s.x, s.y).Length());
                if ((half == 0) == (d > 3.0))
                    keep.push_back(s);
            }
            f.spots = keep;
        }
        const std::string label = half == 0 ? "lowres_d_gt_3A" : "highres_d_lt_3A";
        const Fit f = Place(in.nominal, in.cell, sub, in.start, st, true);
        if (!f.ok) {
            log.Info("ModuleDiag {}: no placement (too few spots)", label);
            continue;
        }
        log.Info("ModuleDiag {}: beam ({:.3f}, {:.3f}) distance {:.4f} spots {}", label, f.r.beam_x_px, f.r.beam_y_px,
                 f.r.distance_mm, f.r.spots_used);
        LogShifts(log, label, in.centres, f.modules, Moved(f.r, st));
    }

    // A wrong start: the beam 3 px off in x and y, the distance 1 mm long.
    {
        DiffractionGeometry g = in.nominal;
        g.BeamX_pxl(g.GetBeamX_pxl() + 3).BeamY_pxl(g.GetBeamY_pxl() + 3).DetectorDistance_mm(g.GetDetectorDistance_mm() + 1);
        const Fit f = Place(g, in.cell, in.frames, in.start, st, true);
        if (f.ok) {
            log.Info("ModuleDiag perturbed_start: beam ({:.3f}, {:.3f}) distance {:.4f}", f.r.beam_x_px, f.r.beam_y_px,
                     f.r.distance_mm);
            LogShifts(log, "perturbed_start", in.centres, f.modules, Moved(f.r, st));
        }
    }

    // The cell free (a weak anchor) instead of pinned.
    {
        GeometryRefinerSettings s2 = st;
        s2.cell_reg_coeff = 0.3;
        const Fit f = Place(in.nominal, in.cell, in.frames, in.start, s2, true);
        if (f.ok) {
            log.Info("ModuleDiag cell_free: beam ({:.3f}, {:.3f}) distance {:.4f} cell {:.3f} {:.3f} {:.3f}",
                     f.r.beam_x_px, f.r.beam_y_px, f.r.distance_mm, f.r.cell.a, f.r.cell.b, f.r.cell.c);
            LogShifts(log, "cell_free", in.centres, f.modules, Moved(f.r, st));
        }
    }

    // The tilt refined as well.
    {
        GeometryRefinerSettings s2 = st;
        s2.refine_tilt = true;
        const Fit f = Place(in.nominal, in.cell, in.frames, in.start, s2, true);
        if (f.ok) {
            log.Info("ModuleDiag tilt_refined: beam ({:.3f}, {:.3f}) distance {:.4f} rot1_deg {:+.4f} rot2_deg {:+.4f} "
                     "median_residual {:.3f}", f.r.beam_x_px, f.r.beam_y_px, f.r.distance_mm, f.r.rot1_rad * RAD2DEG,
                     f.r.rot2_rad * RAD2DEG, f.r.median_residual_px);
            const auto w2 = Moved(f.r, st);
            RadialTable(log, "tilt_refined", f.r.beam_x_px, f.r.beam_y_px, f.modules, w2, f.r.module_mean_d_A,
                        held.r.module_capped_rms_px, f.r.module_capped_rms_px, {});
            Decompose(log, "tilt_refined", f.r.beam_x_px, f.r.beam_y_px, f.modules, w2, lever);
        }
    }

    // Inner against outer modules: the global geometry from each alone, modules held.
    {
        std::vector<double> r;
        for (size_t m: which)
            r.push_back(std::hypot(in.centres[m].x0 - bx, in.centres[m].y0 - by));
        std::vector<double> sorted = r;
        std::sort(sorted.begin(), sorted.end());
        const double median_r = sorted.empty() ? 0 : sorted[sorted.size() / 2];
        for (int outer = 0; outer < 2; outer++) {
            std::vector<GeomRefineFrame> sub = in.frames;
            for (auto &f: sub) {
                std::vector<GeomRefineSpot> keep;
                for (const auto &s: f.spots) {
                    if (s.module < 0)
                        continue;
                    const double rr = std::hypot(in.centres[s.module].x0 - bx, in.centres[s.module].y0 - by);
                    if ((rr >= median_r) == (outer == 1))
                        keep.push_back(s);
                }
                f.spots = keep;
            }
            const Fit f = Place(in.nominal, in.cell, sub, in.centres, st, false);
            log.Info("ModuleDiag {}_modules_only: ok {} beam ({:.3f}, {:.3f}) distance {:.4f} cell {:.3f} {:.3f} {:.3f} "
                     "spots {}", outer ? "outer" : "inner", f.ok, f.r.beam_x_px, f.r.beam_y_px, f.r.distance_mm,
                     f.r.cell.a, f.r.cell.b, f.r.cell.c, f.r.spots_used);
        }
    }

    // Synthetic control: every spot moved onto its prediction at the placed geometry, as found on unmoved
    // modules, with Gaussian noise of its module's rms. Placing these must find no module moves; with a
    // parallax shift added (the mean conversion depth times tan 2theta, radially), whatever it finds is what
    // parallax would put into the modules.
    {
        const double depth_um = in.sensor_thickness_um > 0 ? in.sensor_thickness_um : 300.0;
        const std::string material = in.sensor_material.empty() ? "Si" : in.sensor_material;
        const double L = sensor_absorption::AttenuationLength_um(material, in.nominal.GetWavelength_A());
        // Mean depth of an exponential truncated at the sensor's thickness.
        const double mean_z_um = L - depth_um * std::exp(-depth_um / L) / (1 - std::exp(-depth_um / L));
        const double mean_z_px = mean_z_um / (in.nominal.GetPixelSize_mm() * 1000.0);
        DiffractionGeometry truth = in.nominal;
        truth.BeamX_pxl(static_cast<float>(placed.beam_x_px)).BeamY_pxl(static_cast<float>(placed.beam_y_px))
             .DetectorDistance_mm(static_cast<float>(placed.distance_mm));
        for (int parallax = 0; parallax < 2; parallax++) {
            std::mt19937 noise_rng(777);
            std::vector<GeomRefineFrame> syn = in.frames;
            size_t kept = 0;
            for (size_t i = 0; i < syn.size(); i++) {
                std::vector<GeomRefineSpot> keep;
                for (size_t k = 0; k < syn[i].spots.size(); k++) {
                    auto s = syn[i].spots[k];
                    if (i >= placed.predicted_px.size() || k >= placed.predicted_px[i].size())
                        continue;
                    const auto p = placed.predicted_px[i][k];
                    if (!std::isfinite(p[0]) || s.module < 0)
                        continue;
                    // The prediction is where the placed modules would have the spot found; on modules
                    // with no error it is found where the placed module puts it.
                    const auto &mv = placed.modules[s.module];
                    const double ux = p[0] - mv.x0, uy = p[1] - mv.y0;
                    const double c = std::cos(mv.rot_rad), sn = std::sin(mv.rot_rad);
                    double x = mv.x0 + c * ux - sn * uy + mv.dx;
                    double y = mv.y0 + sn * ux + c * uy + mv.dy;
                    if (parallax) {
                        const double ox = x - truth.GetBeamX_pxl(), oy = y - truth.GetBeamY_pxl();
                        const double rr = std::hypot(ox, oy);
                        const double tan2t = rr * truth.GetPixelSize_mm() / truth.GetDetectorDistance_mm();
                        if (rr > 0) {
                            x += mean_z_px * tan2t * ox / rr;
                            y += mean_z_px * tan2t * oy / rr;
                        }
                    }
                    const double sigma_px = s.module < static_cast<int>(placed.module_capped_rms_px.size())
                                            ? std::max(0.1, placed.module_capped_rms_px[s.module]) : 0.5;
                    std::normal_distribution<double> nd(0.0, sigma_px / std::sqrt(2.0));
                    s.x = static_cast<float>(x + nd(noise_rng));
                    s.y = static_cast<float>(y + nd(noise_rng));
                    keep.push_back(s);
                }
                kept += keep.size();
                syn[i].spots = keep;
            }
            const std::string label = parallax ? "synthetic_parallax" : "synthetic_no_error";
            const Fit f = Place(truth, placed.cell, syn, in.centres, st, true);
            log.Info("ModuleDiag {}: {} spots, mean conversion depth {:.1f} um ({:.3f} px; {} {:.0f} um), ok {}, beam "
                     "({:.3f}, {:.3f}) distance {:.4f}", label, kept, mean_z_um, mean_z_px, material, depth_um, f.ok,
                     f.r.beam_x_px, f.r.beam_y_px, f.r.distance_mm);
            if (f.ok) {
                RadialTable(log, label, f.r.beam_x_px, f.r.beam_y_px, f.modules, Moved(f.r, st), f.r.module_mean_d_A,
                            {}, f.r.module_capped_rms_px, {});
                Decompose(log, label, f.r.beam_x_px, f.r.beam_y_px, f.modules, Moved(f.r, st), lever);
            }
        }
    }
}