// 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) {
#ifdef JFJOCH_USE_CUDA
    std::vector<IndexingResult> ret;

    if (coord.size() <= viable_cell_min_spots)
        return ret;

    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, coord.size());

    fast_feedback::refine::indexer_ifssr<float>::refine(indexer.spotM().topRows(coord.size()),
                                                        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());

    bool indexed = fast_feedback::refine::is_viable_cell(indexer.oCell(id), indexer.Spots(),
                                                         indexing_threshold, viable_cell_min_spots,
                                                         false);

    // Check if result is viable
    if (indexed) {
        auto cell = indexer.oCell(id).colwise().reverse();
        fast_feedback::refine::make_right_handed(cell);

        // get indexed spots
        using M3x = Eigen::MatrixX3<float>;
        M3x resid = indexer.spotM().topRows(coord.size()) * cell.transpose();
        const M3x miller = round(resid.array());
        const M3x predicted = miller * cell.transpose().inverse();

        IndexingResult result;
        result.l = CrystalLattice(cell);

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

        for (int i = 0; i < coord.size(); i++)
            result.predicted_spots[i] = Coord(predicted.coeff(i, 0),
                                              predicted.coeff(i, 1),
                                              predicted.coeff(i, 2));

        ret.emplace_back(result);
    }

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