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

#include "BraggPrediction.h"

std::multimap<float, Reflection> CalcBraggPredictions(const DiffractionExperiment& experiment,
                                                      const CrystalLattice& lattice,
                                                      float high_res_A,
                                                      float ewald_dist_cutoff,
                                                      int max_hkl) {
    std::multimap<float, Reflection> reflections;

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

    for (int h = -max_hkl; h < 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 = -max_hkl; k < 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 = -max_hkl; l < max_hkl; l++) {
                // (000) is not too interesting
                if (h == 0 && k == 0 && l == 0)
                    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 <= ewald_dist_cutoff ) {
                    Coord recip(recip_x, recip_y, recip_z);

                    auto [x,y] = geom.RecipToDector(recip);

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

                    float d = 1.0f / sqrtf(dot);

                    reflections.insert(std::make_pair(dist_ewald_sphere, Reflection{
                            .h = h,
                            .k = k,
                            .l = l,
                            .predicted_x = x,
                            .predicted_y = y,
                            .d = d
                    }));
                }
            }
        }
    }
    return reflections;
}