// 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 "../image_analysis/scale_merge/SphericalHarmonicSurface.h"

namespace {
    constexpr double PI = 3.14159265358979323846;
    // Cell centre direction of the basis grid (the same equal-area layout MakeSphericalHarmonicBasis uses).
    void CellCentre(const SphericalHarmonicBasis &b, int c, double &x, double &y, double &z) {
        const int iz = c / b.nphi, ip = c % b.nphi;
        z = -1.0 + (iz + 0.5) * 2.0 / b.nz;
        const double phi = -PI + (ip + 0.5) * 2.0 * PI / b.nphi;
        const double r = std::sqrt(1.0 - z * z);
        x = r * std::cos(phi);
        y = r * std::sin(phi);
    }

    // Run the fit as the correction-surface engine drives it: every sampled cell sees its observations
    // scaled by the true factor T and the current surface A, so cross = ref2 * T * A, and the fixed
    // point is A = 1/T. Only the band |z| < 0.8 is sampled, as a rotation sweep leaves caps unsampled.
    std::vector<double> Fit(const SphericalHarmonicBasis &b, const std::vector<double> &T) {
        const int ncell = b.NCell();
        std::vector<double> ref2(ncell, 0.0), log_a(ncell, 0.0), theta(b.nterm, 0.0);
        for (int c = 0; c < ncell; ++c) {
            double x, y, z;
            CellCentre(b, c, x, y, z);
            if (std::fabs(z) < 0.8) ref2[c] = 1e4;
        }
        for (int it = 0; it < 10; ++it) {
            std::vector<double> cross(ncell, 0.0);
            for (int c = 0; c < ncell; ++c)
                cross[c] = ref2[c] * T[c] * std::exp(log_a[c]);
            log_a = SphericalHarmonicStep(b, ref2, cross, 0.0, theta);
        }
        return log_a;
    }
}

TEST_CASE("SphericalHarmonics: unit rms and orthogonal over the sphere", "[spherical_harmonics]") {
    const auto b = MakeSphericalHarmonicBasis(128, 256, 6, 0.1);
    REQUIRE(b.nterm == 48);
    const int n = b.NCell();
    for (int i = 0; i < b.nterm; ++i)
        for (int j = i; j < b.nterm; ++j) {
            double s = 0.0;
            for (int c = 0; c < n; ++c)
                s += b.y[static_cast<size_t>(c) * b.nterm + i] * b.y[static_cast<size_t>(c) * b.nterm + j];
            CHECK(s / n == Catch::Approx(i == j ? 1.0 : 0.0).margin(0.02));
        }
}

TEST_CASE("SphericalHarmonicSurface: recovers a smooth absorption surface", "[spherical_harmonics]") {
    const auto b = MakeSphericalHarmonicBasis(32, 64, 6, 0.1);
    std::vector<double> T(b.NCell());
    for (int c = 0; c < b.NCell(); ++c) {
        double x, y, z;
        CellCentre(b, c, x, y, z);
        // A smooth path-length surface: a 15% gradient along one axis plus a 2-fold modulation.
        T[c] = std::exp(0.15 * z + 0.08 * (x * x - y * y) + 0.05 * x * y);
    }
    const auto log_a = Fit(b, T);
    double max_err = 0.0;
    for (int c = 0; c < b.NCell(); ++c) {
        double x, y, z;
        CellCentre(b, c, x, y, z);
        if (std::fabs(z) < 0.8)
            max_err = std::max(max_err, std::fabs(log_a[c] + std::log(T[c])));
    }
    CHECK(max_err < 0.005);
}

TEST_CASE("SphericalHarmonicSurface: inert without an absorption surface", "[spherical_harmonics]") {
    const auto b = MakeSphericalHarmonicBasis(32, 64, 6, 0.1);
    // No surface at all: the coefficients stay at zero.
    const auto flat = Fit(b, std::vector<double>(b.NCell(), 1.0));
    for (double v : flat)
        CHECK(std::fabs(v) < 1e-12);

    // Independent 5% cell-to-cell noise and no smooth structure: the 48 smooth terms take up only
    // a small part of it.
    std::mt19937 rng(7);
    std::normal_distribution<double> noise(0.0, 0.05);
    std::vector<double> T(b.NCell());
    for (double &t : T) t = std::exp(noise(rng));
    const auto log_a = Fit(b, T);
    double s = 0.0;
    int n = 0;
    for (int c = 0; c < b.NCell(); ++c) {
        double x, y, z;
        CellCentre(b, c, x, y, z);
        if (std::fabs(z) < 0.8) { s += log_a[c] * log_a[c]; ++n; }
    }
    CHECK(std::sqrt(s / n) < 0.015);
}