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

#include "BraggPrediction.h"

inline bool odd(const int val) {
    return (val & 1) != 0;
}

BraggPrediction::BraggPrediction(int max_reflections)
    : max_reflections(max_reflections), reflections(max_reflections) {}

const std::vector<Reflection> &BraggPrediction::GetReflections() const {
    return reflections;
}

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

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

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

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

    Coord Astar = lattice.Astar();
    Coord Bstar = lattice.Bstar();
    Coord Cstar = lattice.Cstar();
    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 coeff_const = det_distance / pixel_size;

    int i = 0;

    for (int h = -settings.max_hkl; h < settings.max_hkl; h++) {
        // Precompute A* h contribution
        const float Ah_x = Astar.x * h;
        const float Ah_y = Astar.y * h;
        const float Ah_z = Astar.z * h;

        for (int k = -settings.max_hkl; k < settings.max_hkl; k++) {
            // Accumulate B* k contribution
            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 (int l = -settings.max_hkl; l < settings.max_hkl; l++) {
                bool absent = false;
                // (000) is not too interesting
                if (h == 0 && k == 0 && l == 0)
                    absent = true;

                switch (settings.centering) {
                    case 'I':
                        // Body-centered: h + k + l must be even
                        if (odd(h + k + l))
                            absent = true;
                        break;
                    case 'A':
                        // A-centered: k + l must be even
                        if (odd(k + l))
                            absent = true;
                        break;
                    case 'B':
                        // B-centered: h + l must be even
                        if (odd(h + l))
                            absent = true;
                        break;
                    case 'C':
                        // C-centered: h + k must be even
                        if (odd(h + k))
                            absent = true;
                        break;
                    case 'F': { // Face-centered: h, k, l all even or all odd
                        if (odd(h+k) || odd(h+l) || odd(k+l))
                            absent = true;
                        break;
                    }
                    case 'R': {
                        // Rhombohedral in hexagonal setting (hR, a_h=b_h, γ=120°):
                        // Reflection condition: -h + k + l = 3n  (equivalently h - k + l = 3n)
                        int mod = (-h + k + l) % 3;
                        if (mod < 0) mod += 3;
                        if (mod != 0)
                            absent = true;
                        break;
                    }
                    default:
                        // P or unspecified: no systematic absences
                        break;
                }

                if (absent)
                    continue;

                if (i >= max_reflections)
                    continue;

                float recip_x = AhBk_x + Cstar.x * l;
                float recip_y = AhBk_y + Cstar.y * l;
                float recip_z = AhBk_z + Cstar.z * l;

                float dot = recip_x * recip_x + recip_y * recip_y + recip_z * recip_z;
                if (dot > one_over_dmax_sq)
                    continue;

                float S_x = recip_x + S0.x;
                float S_y = recip_y + S0.y;
                float S_z = recip_z + S0.z;
                float S_len = sqrtf(S_x * S_x + S_y * S_y + S_z * S_z);
                float dist_ewald_sphere = std::fabs(S_len - one_over_wavelength);

                if (dist_ewald_sphere <= settings.ewald_dist_cutoff ) {
                    // 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;

                    // Project to detector coordinates
                    float coeff = coeff_const / S_rot_z;
                    float x = beam_x + S_rot_x * coeff;
                    float y = beam_y + S_rot_y * coeff;

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

                    float d = 1.0f / sqrtf(dot);
                    reflections[i] = Reflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .predicted_x = x,
                        .predicted_y = y,
                        .d = d,
                        .dist_ewald = dist_ewald_sphere
                    };
                    ++i;
                }
            }
        }
    }
    return i;
}