// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "IndexerWrapper.h"

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

// Select spots that belong to indexing solution
// - cell       cell in real space
// - spots      spots in reciprocal space
// - threshold  radius around approximated miller indices
template <typename Mat3, typename MatX3, typename float_type=typename Mat3::Scalar>
inline auto select_indexed_spots(const Eigen::MatrixBase<Mat3>& cell,
                                 const Eigen::MatrixBase<MatX3>& spots,
                                 float_type threshold=.02f)
{
    using M3x = Eigen::MatrixX3<float_type>;
    M3x resid = spots * cell.transpose();
    const M3x miller = round(resid.array());
    resid -= miller;
    return resid.rowwise().norm().array() < threshold;
}

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

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

    size_t nspots = std::min<size_t>(MAX_SPOTS_TO_INDEX, coord.size());

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

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

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

    // Get indexed spots
    auto arr = select_indexed_spots(indexer.oCell(id), indexer.Spots(), threshold);

    auto indexed_spot_count = arr.count();
    // Check if result is viable
    if (indexed_spot_count > min_spots) {
        IndexingResult result;
        result.l = CrystalLattice(indexer.oCell(id));

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

        result.indexed_spots_count = indexed_spot_count;

        ret.emplace_back(result);
    }

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