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

#include "SphericalHarmonicSurface.h"

#include <algorithm>
#include <cmath>
#include <vector>

#include <Eigen/Dense>

namespace {
constexpr double PI = 3.14159265358979323846;
}

void RealSphericalHarmonics(double x, double y, double z, int lmax, double *out) {
    const double ct = std::clamp(z, -1.0, 1.0);
    const double st = std::sqrt(std::max(0.0, 1.0 - ct * ct));
    const double phi = std::atan2(y, x);
    // Associated Legendre P_l^m(cos theta) for 0 <= m <= l <= lmax, by the standard recurrences:
    // P_m^m = (2m-1)!! sin^m, P_{m+1}^m = (2m+1) cos P_m^m,
    // P_l^m = ((2l-1) cos P_{l-1}^m - (l+m-1) P_{l-2}^m) / (l-m). The Condon-Shortley sign is left
    // out; it only flips the sign of a coefficient.
    const int n = lmax + 1;
    std::vector<double> P(n * n, 0.0);
    auto p = [&](int l, int m) -> double & { return P[l * n + m]; };
    p(0, 0) = 1.0;
    for (int m = 1; m <= lmax; ++m)
        p(m, m) = p(m - 1, m - 1) * (2 * m - 1) * st;
    for (int m = 0; m < lmax; ++m)
        p(m + 1, m) = (2 * m + 1) * ct * p(m, m);
    for (int m = 0; m <= lmax; ++m)
        for (int l = m + 2; l <= lmax; ++l)
            p(l, m) = ((2 * l - 1) * ct * p(l - 1, m) - (l + m - 1) * p(l - 2, m)) / (l - m);

    int k = 0;
    for (int l = 1; l <= lmax; ++l)
        for (int m = -l; m <= l; ++m) {
            const int am = std::abs(m);
            // sqrt(4 pi) times the orthonormal normalisation, so each function has unit rms over the sphere.
            double ratio = 1.0;   // (l - |m|)! / (l + |m|)!
            for (int i = l - am + 1; i <= l + am; ++i)
                ratio /= i;
            const double norm = std::sqrt((2 * l + 1) * ratio);
            if (m == 0)
                out[k++] = norm * p(l, 0);
            else if (m > 0)
                out[k++] = std::sqrt(2.0) * norm * p(l, am) * std::cos(am * phi);
            else
                out[k++] = std::sqrt(2.0) * norm * p(l, am) * std::sin(am * phi);
        }
}

SphericalHarmonicBasis MakeSphericalHarmonicBasis(int nz, int nphi, int lmax, double prior_sigma) {
    SphericalHarmonicBasis b;
    b.lmax = lmax;
    b.nterm = (lmax + 1) * (lmax + 1) - 1;
    b.nz = nz;
    b.nphi = nphi;
    b.y.resize(static_cast<size_t>(b.NCell()) * b.nterm);
    // Bands equal in u.z are equal in area, so every cell covers the same solid angle.
    for (int iz = 0; iz < nz; ++iz)
        for (int ip = 0; ip < nphi; ++ip) {
            const double z = -1.0 + (iz + 0.5) * 2.0 / nz;
            const double phi = -PI + (ip + 0.5) * 2.0 * PI / nphi;
            const double r = std::sqrt(std::max(0.0, 1.0 - z * z));
            RealSphericalHarmonics(r * std::cos(phi), r * std::sin(phi), z, lmax,
                                   b.y.data() + static_cast<size_t>(iz * nphi + ip) * b.nterm);
        }
    b.prior.resize(b.nterm);
    for (int l = 1, k = 0; l <= lmax; ++l)
        for (int m = -l; m <= l; ++m, ++k) {
            const double s = prior_sigma / l;
            b.prior[k] = 1.0 / (s * s);
        }
    return b;
}

int SphericalHarmonicCell(const SphericalHarmonicBasis &basis, double x, double y, double z) {
    const int iz = std::clamp(static_cast<int>((z + 1.0) * 0.5 * basis.nz), 0, basis.nz - 1);
    const int ip = std::clamp(static_cast<int>((std::atan2(y, x) + PI) / (2.0 * PI) * basis.nphi),
                              0, basis.nphi - 1);
    return iz * basis.nphi + ip;
}

std::vector<double> SphericalHarmonicStep(const SphericalHarmonicBasis &basis, const std::vector<double> &ref2,
                                          const std::vector<double> &cross, double damping,
                                          std::vector<double> &theta) {
    const int K = basis.nterm, ncell = basis.NCell();
    Eigen::MatrixXd H = Eigen::MatrixXd::Zero(K, K);
    Eigen::VectorXd g = Eigen::VectorXd::Zero(K);
    for (int c = 0; c < ncell; ++c) {
        if (!(ref2[c] > 0.0)) continue;
        const Eigen::Map<const Eigen::VectorXd> yc(basis.y.data() + static_cast<size_t>(c) * K, K);
        H.noalias() += (ref2[c] + damping) * yc * yc.transpose();
        g += (ref2[c] - cross[c]) * yc;
    }
    for (int k = 0; k < K; ++k) {
        H(k, k) += basis.prior[k];
        g(k) -= basis.prior[k] * theta[k];
    }
    const Eigen::VectorXd step = H.ldlt().solve(g);
    for (int k = 0; k < K; ++k)
        theta[k] += step(k);

    std::vector<double> log_a(ncell);
    const Eigen::Map<const Eigen::VectorXd> th(theta.data(), K);
    for (int c = 0; c < ncell; ++c)
        log_a[c] = Eigen::Map<const Eigen::VectorXd>(basis.y.data() + static_cast<size_t>(c) * K, K).dot(th);
    return log_a;
}