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

#include <cstdint>
#include <vector>
#include "AnalyzeIndexing.h"

#include "FitProfileRadius.h"

inline bool ok(float x) {
    if (!std::isfinite(x))
        return false;
    if (x < 0.0)
        return false;
    return true;
};

bool AnalyzeIndexing(DataMessage &message, const DiffractionExperiment &experiment, const CrystalLattice &latt) {

    size_t nspots = message.spots.size();

    uint64_t indexed_spot_count = 0;
    std::vector<uint8_t> indexed_spots(nspots);

    std::vector<float> distance_ewald_sphere(nspots);

    // Check spots
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    const Coord astar = latt.Astar();
    const Coord bstar = latt.Bstar();
    const Coord cstar = latt.Cstar();

    const auto geom = experiment.GetDiffractionGeometry();
    const auto indexing_tolerance = experiment.GetIndexingSettings().GetTolerance();
    const auto viable_cell_min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots();
     // identify indexed spots
    for (int i = 0; i < message.spots.size(); i++) {
        auto recip = message.spots[i].ReciprocalCoord(geom);

        float h_fp = recip * a;
        float k_fp = recip * b;
        float l_fp = recip * c;

        float h_frac = h_fp - std::round(h_fp);
        float k_frac = k_fp - std::round(k_fp);
        float l_frac = l_fp - std::round(l_fp);

        float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac;

        Coord recip_pred = std::round(h_fp) * astar + std::round(k_fp) * bstar + std::round(l_fp) * cstar;

        // See indexing_peak_check() in peaks.c in CrystFEL
        if (norm_sq < indexing_tolerance * indexing_tolerance) {
            indexed_spot_count++;
            indexed_spots[i] = 1;
            message.spots[i].dist_ewald_sphere = geom.DistFromEwaldSphere(recip_pred);
            message.spots[i].h = std::lround(h_fp);
            message.spots[i].k = std::lround(k_fp);
            message.spots[i].l = std::lround(l_fp);
        }
    }

    auto spot_count_threshold = std::max<int64_t>(viable_cell_min_spots, std::lround(min_percentage_spots * nspots));

    if (indexed_spot_count >= spot_count_threshold) {
        auto uc = latt.GetUnitCell();
        if (!ok(uc.a) || !ok(uc.b) || !ok(uc.c) || !ok(uc.alpha) || !ok(uc.beta) || !ok(uc.gamma))
            return {};

        message.indexing_result = true;
        assert(indexed_spots.size() == message.spots.size());
        for (int i = 0; i < message.spots.size(); i++)
            message.spots[i].indexed = indexed_spots[i];

        message.profile_radius = FitProfileRadius(message.spots);
        message.spot_count_indexed = indexed_spot_count;
        message.indexing_lattice = latt;
        message.indexing_unit_cell = latt.GetUnitCell();
        return true;
    }
    message.indexing_result = false;
    return false;
}