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

#include "BraggPredictionRotation.h"

#include <cmath>
#include <algorithm>

#include "../../common/DiffractionGeometry.h"
#include "../../common/GoniometerAxis.h"
#include "../../common/JFJochException.h"
#include "SystematicAbsence.h"

namespace {
    inline float deg_to_rad(float deg) {
        return deg * (static_cast<float>(M_PI) / 180.0f);
    }
    inline float rad_to_deg(float rad) {
        return rad * (180.0f / static_cast<float>(M_PI));
    }

    // Solve A cos(phi) + B sin(phi) + D = 0, return solutions in [phi0, phi1]
    // Returns 0..2 solutions.
    inline int solve_trig(float A, float B, float D,
                          float phi0, float phi1,
                          float out_phi[2]) {
        const float R = std::sqrt(A*A + B*B);
        if (!(R > 0.0f))
            return 0;

        const float rhs = -D / R;
        if (rhs < -1.0f || rhs > 1.0f)
            return 0;

        const float phi_ref = std::atan2(B, A);
        const float delta = std::acos(rhs);

        float s1 = phi_ref + delta;
        float s2 = phi_ref - delta;

        // Shift candidates by +/- 2*pi so they land near the interval.
        const float two_pi = 2.0f * static_cast<float>(M_PI);
        auto shift_near = [&](float x) {
            while (x < phi0 - two_pi) x += two_pi;
            while (x > phi1 + two_pi) x -= two_pi;
            return x;
        };

        s1 = shift_near(s1);
        s2 = shift_near(s2);

        int n = 0;
        if (s1 >= phi0 && s1 <= phi1) out_phi[n++] = s1;
        if (s2 >= phi0 && s2 <= phi1) {
            if (n == 0 || std::fabs(s2 - out_phi[0]) > 1e-6f) out_phi[n++] = s2;
        }
        return n;
    }

    // Convert angle to fractional image_number using goniometer start/increment
    inline float phi_deg_to_image_number(float phi_deg, float start_deg, float increment_deg) {
        return (phi_deg - start_deg) / increment_deg;
    }
} // namespace

std::vector<Reflection> BraggPredictionRotation::Calc(const DiffractionExperiment& experiment,
                                                        const CrystalLattice& lattice,
                                                        const BraggPredictionRotationSettings& settings) const {
    std::vector<Reflection> out;
    out.reserve(200000); // tune

    const auto gon_opt = experiment.GetGoniometer();
    if (!gon_opt.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "BraggPredictionRotationCPU requires a goniometer axis");

    const GoniometerAxis& gon = *gon_opt;
    const auto geom = experiment.GetDiffractionGeometry();

    // Determine prediction interval in spindle angle.
    // We treat each image as [angle(image), angle(image)+wedge)
    const float start_deg = gon.GetStart_deg();
    const float inc_deg   = gon.GetIncrement_deg();
    const float wedge_deg = gon.GetWedge_deg();

    float phi0_deg = gon.GetAngle_deg(static_cast<float>(settings.image_first));
    float phi1_deg = gon.GetAngle_deg(static_cast<float>(settings.image_last)) + wedge_deg;

    if (settings.pad_one_wedge) {
        phi0_deg -= wedge_deg;
        phi1_deg += wedge_deg;
    }

    const float phi0 = deg_to_rad(phi0_deg);
    const float phi1 = deg_to_rad(phi1_deg);

    // Resolution cutoff in reciprocal space
    const float one_over_dmax = 1.0f / settings.high_res_A;
    const float one_over_dmax_sq = one_over_dmax * one_over_dmax;

    // S0 has length 1/lambda (your convention)
    const Coord S0 = geom.GetScatteringVector();
    const float k2 = (S0 * S0); // (1/lambda)^2

    // Rotation axis in lab frame (already normalized in ctor, but keep safe)
    const Coord w = gon.GetAxis().Normalize();

    const Coord Astar = lattice.Astar();
    const Coord Bstar = lattice.Bstar();
    const Coord Cstar = lattice.Cstar();

    const float det_w = static_cast<float>(experiment.GetXPixelsNum());
    const float det_h = static_cast<float>(experiment.GetYPixelsNum());

    for (int h = -settings.max_hkl; h <= settings.max_hkl; ++h) {
        const Coord Ah = Astar * static_cast<float>(h);

        for (int k = -settings.max_hkl; k <= settings.max_hkl; ++k) {
            const Coord AhBk = Ah + Bstar * static_cast<float>(k);

            for (int l = -settings.max_hkl; l <= settings.max_hkl; ++l) {
                if (systematic_absence(h, k, l, settings.centering))
                    continue;

                const Coord g0 = AhBk + Cstar * static_cast<float>(l);
                const float g02 = g0 * g0;

                if (!(g02 > 0.0f))
                    continue;
                if (g02 > one_over_dmax_sq)
                    continue;

                // g0 = g_par + g_perp wrt spindle axis
                const float g_par_s = g0 * w;
                const Coord g_par = w * g_par_s;
                const Coord g_perp = g0 - g_par;

                const float g_perp2 = g_perp * g_perp;
                if (g_perp2 < 1e-12f) {
                    // Rotation does not move this reciprocal vector: skip for now.
                    // (Could be handled by checking whether it satisfies Ewald at all.)
                    continue;
                }

                // Equation: |S0 + g(phi)|^2 = |S0|^2
                // g(phi) = g_par + cos(phi) g_perp + sin(phi) (w x g_perp)
                const Coord p = S0 + g_par;
                const Coord w_x_gperp = w % g_perp;

                const float A = 2.0f * (p * g_perp);
                const float B = 2.0f * (p * w_x_gperp);
                const float D = (p * p) + g_perp2 - k2;

                float sols[2]{};
                const int nsol = solve_trig(A, B, D, phi0, phi1, sols);
                if (nsol == 0)
                    continue;

                for (int si = 0; si < nsol; ++si) {
                    const float phi = sols[si];

                    // Rotate g0 by phi about w
                    const RotMatrix R(phi, w);
                    const Coord g = R * g0;
                    const Coord S = S0 + g;

                    // Project to detector using canonical geometry code
                    const auto [x, y] = geom.RecipToDector(g);
                    if (!std::isfinite(x) || !std::isfinite(y))
                        continue;
                    if (x < 0.0f || x >= det_w || y < 0.0f || y >= det_h)
                        continue;

                    // Convert phi to fractional image number
                    const float phi_deg = rad_to_deg(phi);

                    Reflection r{};
                    r.h = h;
                    r.k = k;
                    r.l = l;
                    r.angle_deg = phi_deg;
                    r.image_number = phi_deg_to_image_number(phi_deg, start_deg, inc_deg);
                    r.predicted_x = x;
                    r.predicted_y = y;
                    r.d = 1.0f / std::sqrt(g02);

                    // diagnostic: should be ~0
                    r.dist_ewald = S.Length() - std::sqrt(k2);

                    r.I = 0.0f;
                    r.bkg = 0.0f;
                    r.sigma = 0.0f;

                    out.push_back(r);
                }
            }
        }
    }

    std::sort(out.begin(), out.end(), [](const Reflection &a, const Reflection &b) {
        if (a.angle_deg != b.angle_deg)
            return a.angle_deg < b.angle_deg;
        if (a.h != b.h) return a.h < b.h;
        if (a.k != b.k) return a.k < b.k;
        return a.l < b.l;
    });

    return out;
}