// Copyright (2019-2024) Paul Scherrer Institute

#include "IndexerWrapper.h"

void IndexerWrapper::Setup(const UnitCell &cell) {
#ifdef JFJOCH_USE_CUDA
    indexer.iCellM() = CrystalLattice(cell).GetEigenMatrix();
#endif
}

std::vector<IndexingResult> IndexerWrapper::Run(const std::vector<Coord> &coord, float indexing_threshold, int nspots) {
#ifdef JFJOCH_USE_CUDA
    std::vector<IndexingResult> ret;

    if (nspots > coord.size())
        nspots = coord.size();

    if (nspots <= viable_cell_min_spots)
        return ret;

    assert(nspots <= MAX_SPOT_COUNT);
    assert(coord.size() <= MAX_SPOT_COUNT);

    for (int i = 0; i < coord.size(); i++) {
        indexer.spotX(i) = coord[i].x;
        indexer.spotY(i) = coord[i].y;
        indexer.spotZ(i) = coord[i].z;
    }

    // Index
    indexer.index(1, nspots);

    fast_feedback::refine::indexer_ifssr<float>::refine(indexer.spotM().topRows(nspots),
                                                        indexer.oCellM(),
                                                        indexer.oScoreV(),
                                                        cifssr);

    // Find cell most similar to the current one
    for (unsigned i=0u; i<cpers.max_output_cells; i++)
        indexer.oScore(i) += fast_feedback::refine::cell_similarity(indexer.oCell(i), indexer.iCell(0), 0.02f);

    // Get best cell
    auto id = fast_feedback::refine::best_cell(indexer.oScoreV());

    auto cell = indexer.oCell(id).colwise().reverse();
    fast_feedback::refine::make_right_handed(cell);

    using M3x = Eigen::MatrixX3<float>;
    M3x resid = indexer.spotM().topRows(coord.size()) * cell.transpose();
    const M3x miller = round(resid.array());
    resid -= miller;

    auto indexed_spots = (resid.rowwise().norm().array() < indexing_threshold);
    auto indexed_spot_count = indexed_spots.topRows(nspots).count();

    // Conditions for indexing:
    // * There must be 9 AND 20% of all spots indexed
    // * Decision is only based on spots used for indexing (not all spots)
    if ((indexed_spot_count > viable_cell_min_spots) && (indexed_spot_count > 0.20 * nspots)) {
        IndexingResult result;
        result.l = CrystalLattice(cell);

        result.indexed_spot.resize(coord.size());

        for (int i = 0; i < coord.size(); i++)
            result.indexed_spot[i] = indexed_spots(i);

        ret.emplace_back(result);
    }

    return ret;
#else
    return {};
#endif
}