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

#include <array>
#include <cmath>
#include <vector>

#include "SpindleCuspLoss.h"
#include "../common/JFJochMath.h"

std::optional<SpindleCuspLoss> SpindleUnrepairedFraction(const gemmi::SpaceGroup &sg,
                                                         const CrystalLattice &lattice,
                                                         const Coord &spindle,
                                                         double wavelength_A,
                                                         double d_min_A) {
    const double spindle_len = spindle.Length();
    if (spindle_len < 1e-9 || wavelength_A <= 0 || d_min_A <= 0)
        return {};
    const Coord axis = spindle * static_cast<float>(1.0 / spindle_len);

    const double sin_tm = std::min(1.0, wavelength_A / (2.0 * d_min_A));
    const double cos_tm = std::sqrt(std::max(0.0, 1.0 - sin_tm * sin_tm));
    if (sin_tm <= 0)
        return {};

    // Proper rotations only, identity included. Friedel and the improper operators need no separate
    // treatment: the blind double cone and the measured region are both inversion-symmetric, so -R
    // lands a point inside or outside exactly as R does.
    std::vector<std::array<std::array<double, 3>, 3>> rot;
    for (const auto &op : sg.operations().derive_symmorphic().sym_ops) {
        const auto &w = op.rot;
        const double det =
                static_cast<double>(w[0][0]) * (w[1][1] * w[2][2] - w[1][2] * w[2][1])
              - static_cast<double>(w[0][1]) * (w[1][0] * w[2][2] - w[1][2] * w[2][0])
              + static_cast<double>(w[0][2]) * (w[1][0] * w[2][1] - w[1][1] * w[2][0]);
        if (det <= 0)
            continue;
        if (op.rot == gemmi::Op::identity().rot)
            continue;   // handled exactly below, without the float round trip through the basis
        std::array<std::array<double, 3>, 3> m{};
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                m[i][j] = static_cast<double>(w[i][j]) / gemmi::Op::DEN;
        rot.push_back(m);
    }

    const Coord a = lattice.Vec0(), b = lattice.Vec1(), c = lattice.Vec2();
    const Coord as = lattice.Astar(), bs = lattice.Bstar(), cs = lattice.Cstar();

    // Deterministic spiral over ONE lobe of the blind cone; the other lobe contributes identically
    // (the orbit of -p is the negated orbit of p, and only |cos| to the spindle enters). Only cone
    // directions can be lost - the identity is in every orbit - so nothing outside is sampled.
    constexpr int N_DIRECTIONS = 8192;
    const double golden_angle = PI * (3.0 - std::sqrt(5.0));

    // Any unit vector perpendicular to the spindle, to open the cap around it.
    const Coord seed = std::fabs(axis.x) < 0.9f ? Coord(1, 0, 0) : Coord(0, 1, 0);
    const Coord e1 = (axis % seed).Normalize();
    const Coord e2 = axis % e1;

    double lost_sum = 0, p1_sum = 0;
    for (int i = 0; i < N_DIRECTIONS; i++) {
        const double z = cos_tm + (1.0 - cos_tm) * (static_cast<double>(i) + 0.5) / N_DIRECTIONS;
        const double r = std::sqrt(std::max(0.0, 1.0 - z * z));
        const double phi = golden_angle * i;
        const Coord p = axis * static_cast<float>(z)
                      + e1 * static_cast<float>(r * std::cos(phi))
                      + e2 * static_cast<float>(r * std::sin(phi));

        // h = (a.p, b.p, c.p) are continuous Miller coordinates; an operator sends the reflection
        // h to h' = h W (gemmi's row convention) and h' returns to Cartesian through the reciprocal
        // basis. No matrix inversion: direct and reciprocal bases are mutually inverse.
        const double h0 = a * p, h1 = b * p, h2 = c * p;

        // The orbit's largest folded angle to the spindle, as its sine (the fold to [0, 90] makes
        // the sine monotone). A point at direction p and radius q is unmeasured iff every image of
        // its orbit is inside the cone of ITS OWN shell, i.e. iff q > 2 sin(a_max) / lambda - so
        // the direction loses the radial fraction 1 - (sin a_max / sin theta_max)^3 of its unique
        // reflections, reciprocal volume being what unique reflections are proportional to.
        // The identity's own angle comes straight from z, exactly; only the non-trivial
        // operators go through the basis.
        const double sin_id = std::sqrt(std::max(0.0, 1.0 - z * z));
        double max_sin = sin_id;
        for (const auto &m : rot) {
            const double g0 = h0 * m[0][0] + h1 * m[1][0] + h2 * m[2][0];
            const double g1 = h0 * m[0][1] + h1 * m[1][1] + h2 * m[2][1];
            const double g2 = h0 * m[0][2] + h1 * m[1][2] + h2 * m[2][2];
            const Coord s = as * static_cast<float>(g0) + bs * static_cast<float>(g1)
                          + cs * static_cast<float>(g2);
            const double len = s.Length();
            if (len < 1e-12)
                continue;
            const double cos_a = std::min(1.0, std::fabs(s * axis) / len);
            max_sin = std::max(max_sin, std::sqrt(std::max(0.0, 1.0 - cos_a * cos_a)));
        }
        if (max_sin < sin_tm) {
            const double x = max_sin / sin_tm;
            lost_sum += 1.0 - x * x * x;
        }
        const double y = sin_id / sin_tm;
        p1_sum += 1.0 - y * y * y;
    }

    SpindleCuspLoss ret;
    ret.theta_max_deg = std::asin(sin_tm) * 180.0 / PI;
    // The double cone is (1 - cos theta_max) of the sphere, and the sample covers exactly the cone,
    // so the cone average rescales by that solid angle to a fraction of ALL unique reflections.
    ret.lost_unique_fraction = (1.0 - cos_tm) * lost_sum / N_DIRECTIONS;
    ret.cone_fraction = p1_sum > 0 ? lost_sum / p1_sum : 0.0;
    return ret;
}