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

#include "Indexer.h"

void Indexer::Setup(const DiffractionExperiment& experiment) {
    geom = experiment.GetDiffractionGeometry();
    indexing_tolerance = experiment.GetIndexingSettings().GetTolerance();
    indexing_tolerance_sq = indexing_tolerance * indexing_tolerance;
    SetupUnitCell(experiment.GetUnitCell());
}

std::optional<CrystalLattice>
Indexer::Run(DataMessage &message, const std::vector<DiffractionSpot> &spots) {
    auto start = std::chrono::high_resolution_clock::now();
    size_t nspots = spots.size();

    std::vector<Coord> recip;
    recip.reserve(spots.size());

    for (const auto &i: spots)
        recip.push_back(i.ReciprocalCoord(geom));

    auto ret = Run(recip, nspots);

    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<float> duration = end - start;
    message.indexing_time_s = duration.count();

    if (ret.empty()) {
        message.indexing_result = false;
        return {};
    }

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

    // Check spots
    Coord a = ret[0].Vec0();
    Coord b = ret[0].Vec1();
    Coord c = ret[0].Vec2();

    Coord astar = ret[0].Astar();
    Coord bstar = ret[0].Bstar();
    Coord cstar = ret[0].Cstar();

     // identify indexed spots
    for (int i = 0; i < spots.size(); i++) {
        float h_fp = recip[i] * a;
        float k_fp = recip[i] * b;
        float l_fp = recip[i] * c;

        float h = std::round(h_fp);
        float k = std::round(k_fp);
        float l = std::round(l_fp);

        float h_frac = h_fp - h;
        float k_frac = k_fp - k;
        float l_frac = l_fp - l;

        Coord pred_recip = h * astar + k * bstar + l * cstar;
        float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac;

        if (norm_sq < indexing_tolerance_sq) {
            indexed_spot_count++;
            indexed_spots[i] = 1;

            // Assume that angle mean is zero
            float tmp = geom.AngleFromEwaldSphere(pred_recip);
            angle_rms += tmp * tmp;
        }
    }

    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) {
        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.indexing_lattice = ret[0];
        message.indexing_unit_cell = ret[0].GetUnitCell();
        message.mosaicity = sqrtf(angle_rms / static_cast<float>(indexed_spot_count));
        return ret[0];
    }
    message.indexing_result = false;
    return {};
}