// 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 "../bragg_integration/SystematicAbsence.h"

std::vector<PredictionRotationResult> PredictRotationHKLs(const DiffractionExperiment &experiment,
                 const CrystalLattice &lattice,
                 const PredictionRotationSettings &settings) {
    std::vector<PredictionRotationResult> ret;

    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();

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

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

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

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

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

    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;

    for (int32_t h = -settings.max_hkl; h <= settings.max_hkl; ++h) {
        const float Ah_x = Astar.x * h;
        const float Ah_y = Astar.y * h;
        const float Ah_z = Astar.z * h;

        for (int32_t k = -settings.max_hkl; k <= settings.max_hkl; ++k) {
            const float AhBk_x = Ah_x + Bstar.x * k;
            const float AhBk_y = Ah_y + Bstar.y * k;
            const float AhBk_z = Ah_z + Bstar.z * k;

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

                const float p0_x = AhBk_x + Cstar.x * l;
                const float p0_y = AhBk_y + Cstar.y * l;
                const float p0_z = AhBk_z + Cstar.z * l;
                const Coord p0{p0_x, p0_y, p0_z};

                const 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) {
                    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 S = S0 + m1 * p_m1 + m2 * p_m2 + m3 * p_m3;

                    float phi = -1.0f * std::atan2(sinphi, cosphi) * 180.0f / M_PI;
                    if (phi < 0.0f) phi += 360.0f;
                    const float lorentz_reciprocal = std::fabs(m2 * (S % S0))/(S*S0);

                    // 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;

                    ret.emplace_back(PredictionRotationResult{
                        .angle_deg = phi,
                        .lorentz_reciprocal = lorentz_reciprocal,
                        .h = h,
                        .k = k,
                        .l = l,
                        .x = x,
                        .y = y
                    });
                }
            }
        }
    }

    std::ranges::sort(ret,
                      [](const PredictionRotationResult& a, const PredictionRotationResult& b) {
                          return a.angle_deg < b.angle_deg;
                      });
    return ret;
}