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

#include "BraggPredictionRot.h"
#include "../bragg_integration/SystematicAbsence.h"


int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice,
    const BraggPredictionSettings &settings) {

    const auto geom = experiment.GetDiffractionGeometry();
    const auto det_width_pxl = static_cast<float>(experiment.GetXPixelsNum());
    const auto det_height_pxl = static_cast<float>(experiment.GetYPixelsNum());

    const float one_over_dmax = 1.0f / settings.high_res_A;
    const float one_over_dmax_sq = one_over_dmax * one_over_dmax;

    float one_over_wavelength = 1.0f / geom.GetWavelength_A();

    const Coord Astar = lattice.Astar();
    const Coord Bstar = lattice.Bstar();
    const Coord Cstar = lattice.Cstar();
    const Coord S0 = geom.GetScatteringVector();

    std::vector<float> rot = geom.GetPoniRotMatrix().transpose().arr();

    // Precompute detector geometry constants
    float beam_x = geom.GetBeamX_pxl();
    float beam_y = geom.GetBeamY_pxl();
    float det_distance = geom.GetDetectorDistance_mm();
    float pixel_size = geom.GetPixelSize_mm();
    float F = det_distance / pixel_size;

    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 Coord m2 = gon.GetAxis().Normalize();
    const Coord m1 = (m2 % S0).Normalize();
    const Coord m3 = (m1 % m2).Normalize();

    const float m2_S0 = m2 * S0;
    const float m3_S0 = m3 * S0;

    int i = 0;

    const float mos_angle_rad = settings.mosaicity_deg * static_cast<float>(M_PI) / 180.f;
    const float half_wedge_angle_rad = settings.wedge_deg * static_cast<float>(M_PI) / 180.f / 2.0f ;

    for (int h = -settings.max_hkl; h <= settings.max_hkl; h++) {
        // Precompute A* h contribution

        for (int k = -settings.max_hkl; k <= settings.max_hkl; k++) {
            // Accumulate B* k contribution

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

                if (i >= max_reflections)
                    continue;
                Coord p0 = Astar * h + Bstar * k + Cstar * l;

                float p0_sq = p0 * p0;
                if (p0_sq <= 0.0f || p0_sq > one_over_dmax_sq)
                    continue;

                const float p0_m1 = p0 * m1;
                const float p0_m2 = p0 * m2;
                const float p0_m3 = p0 * m3;

                const float rho_sq = p0_sq - (p0_m2 * p0_m2);

                const float p_m3 = (- p0_sq / 2 - p0_m2 * m2_S0) / m3_S0;
                const float p_m2 = p0_m2;
                const float p_m1_opt[2] = {
                    std::sqrt(rho_sq - p_m3 * p_m3),
                    -std::sqrt(rho_sq - p_m3 * p_m3)
                };

                // No solution for Laue equations
                if ((rho_sq < p_m3 * p_m3) || (p0_sq > 4 * S0 * S0))
                    continue;

                for (const auto& p_m1 : p_m1_opt) {
                    if (i >= max_reflections)
                        continue;

                    const float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq;
                    const float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq;
                    Coord p = m1 * p_m1 + m2 * p_m2 + m3 * p_m3; // p0 vector "rotated" to diffracting condition
                    Coord S = S0 + p;

                    float phi = -1.0f * std::atan2(sinphi, cosphi);

                    const Coord e1 = (S % S0).Normalize();

                    const float zeta_abs = std::fabs(m2 * e1);

                    if (zeta_abs < settings.min_zeta)
                        continue;

                    float epsilon3 = std::fabs(phi * zeta_abs);

                    if (epsilon3 > settings.mosaicity_multiplier * mos_angle_rad)
                        continue;

                    const float lorentz_reciprocal = std::fabs(m2 * (S % S0))/(S*S0);
                    const float c1 = zeta_abs / (std::sqrt(2.0f) * mos_angle_rad);

                    const float partiality = (std::erf((phi + half_wedge_angle_rad) * c1)
                        - std::erf((phi - half_wedge_angle_rad) * c1)) / 2.0f;
                    // Inlined RecipToDector with rot1 and rot2 (rot3 = 0)
                    // Apply rotation matrix transpose
                    float S_rot_x = rot[0] * S.x + rot[1] * S.y + rot[2] * S.z;
                    float S_rot_y = rot[3] * S.x + rot[4] * S.y + rot[5] * S.z;
                    float S_rot_z = rot[6] * S.x + rot[7] * S.y + rot[8] * S.z;

                    if (S_rot_z <= 0)
                        continue;

                    float x = beam_x + F * S_rot_x / S_rot_z;
                    float y = beam_y + F * S_rot_y / S_rot_z;

                    if ((x < 0) || (x >= det_width_pxl) || (y < 0) || (y >= det_height_pxl))
                        continue;

                    float dist_ewald_sphere = std::fabs(S.Length() - one_over_wavelength);

                    float d = 1.0f / sqrtf(p0_sq);
                    reflections[i] = Reflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .delta_phi_deg = phi * 180.0f / static_cast<float>(M_PI),
                        .predicted_x = x,
                        .predicted_y = y,
                        .d = d,
                        .dist_ewald = dist_ewald_sphere,
                        .rlp = lorentz_reciprocal,
                        .partiality = partiality,
                        .zeta = zeta_abs
                    };
                    i++;
                }
            }
        }
    }
    return i;
}