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

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

    for (int h = -max_hkl; h < max_hkl; h++) {
        for (int k = -max_hkl; k < max_hkl; k++) {
            for (int l = -max_hkl; l < max_hkl; l++) {
                Coord recip = Astar * static_cast<float>(h)
                              + Bstar * static_cast<float>(k)
                              + Cstar * static_cast<float>(l);
                float d = 1.0f / recip.Length();
                if (d < high_res_A)
                    continue;

                float dist_ewald_sphere = geom.DistFromEwaldSphere(recip);
                float angle_ewald_sphere = geom.AngleFromEwaldSphere_deg(recip);
                if (!std::isfinite(angle_ewald_sphere) || angle_ewald_sphere > angle_dist_cutoff )
                    continue;

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

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

                reflections.insert(std::make_pair(angle_ewald_sphere, Reflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .center_x_pxl = x,
                        .center_y_pxl = y,
                        .radius_ewald_sphere = dist_ewald_sphere,
                        .angle_ewald_sphere = angle_ewald_sphere,
                        .d = d
                }));
            }
        }
    }
    return reflections;
}