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

#include "IndexerWrapper.h"

#ifdef JFJOCH_USE_CUDA
#include <ffbidx/refine.h>
#endif

struct IndexerWrapperImpl {
#ifdef JFJOCH_USE_CUDA
    fast_feedback::config_runtime<float> crt{};
    fast_feedback::config_persistent<float> cpers{
            .max_output_cells = 32,
            .max_input_cells = 1,
            .max_spots = MAX_SPOT_COUNT,
            .num_candidate_vectors = 32,
            .redundant_computations = true,
    };    // default persistent config
    fast_feedback::refine::config_ifssr<float> cifssr{
            .min_spots = 6u
    };
    fast_feedback::refine::indexer<float> indexer{cpers, crt};
#else
    int dummy_variable_to_avoid_empty_struct = 0;
#endif
};

void IndexerWrapper::Setup(const UnitCell &cell) {
#ifdef JFJOCH_USE_CUDA
    CrystalLattice l1(cell);

    Eigen::Matrix3f m;
    impl_->indexer.iCellM() << l1.Vec0().x, l1.Vec0().y, l1.Vec0().z,
            l1.Vec1().x, l1.Vec1().y, l1.Vec1().z,
            l1.Vec2().x, l1.Vec2().y, l1.Vec2().z;
#endif
}

std::vector<IndexingResult> IndexerWrapper::Run(const std::vector<Coord> &coord, float indexing_threshold, size_t 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++) {
        impl_->indexer.spotX(i) = coord[i].x;
        impl_->indexer.spotY(i) = coord[i].y;
        impl_->indexer.spotZ(i) = coord[i].z;
    }

    // Index
    impl_->indexer.index(1, nspots);

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

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

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

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

    using M3x = Eigen::MatrixX3<float>;
    M3x resid = impl_->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.Vec0().x = cell(0,0);
        result.l.Vec0().y = cell(0,1);
        result.l.Vec0().z = cell(0,2);
        result.l.Vec1().x = cell(1,0);
        result.l.Vec1().y = cell(1,1);
        result.l.Vec1().z = cell(1,2);
        result.l.Vec2().x = cell(2,0);
        result.l.Vec2().y = cell(2,1);
        result.l.Vec2().z = cell(2,2);

        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
}

void IndexerWrapper::Run(DataMessage &message,
                         const std::vector<DiffractionSpot> &spots,
                         const DiffractionGeometry &geom,
                         float indexing_threshold) {

    size_t nspots = spots.size();

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

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

    std::vector<IndexingResult> ret;

    if (nspots > 80)
        ret = Run(recip, indexing_threshold, 50);
    if (ret.empty())
        ret = Run(recip, indexing_threshold, nspots);

    if (!ret.empty()) {
        message.indexing_result = true;
        assert(ret[0].indexed_spot.size() == message.spots.size());

        // identify indexed spots
        for (int i = 0; i < message.spots.size(); i++)
            message.spots[i].indexed = ret[0].indexed_spot[i];

        message.indexing_lattice = ret[0].l.GetVector();
        message.indexing_unit_cell = ret[0].l.GetUnitCell();
    }
}

IndexerWrapper::IndexerWrapper() {
    impl_ = std::make_unique<IndexerWrapperImpl>();
}

IndexerWrapper::~IndexerWrapper() {
    impl_.reset();
}