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

#include "FFTIndexer.h"
#include <ffbidx/refine.h>

std::vector<CrystalLattice> FFTIndexer::Refine(const std::vector<Coord> &in_spots, size_t nspots,
                                               const std::vector<CrystalLattice> &lattices) const {
    using M3x = Eigen::MatrixX3<float>;

    std::vector<CrystalLattice> ret;

    Eigen::MatrixX3<float> spots(in_spots.size(), 3u);
    Eigen::MatrixX3<float> oCell(lattices.size() * 3u, 3u);
    Eigen::VectorX<float> scores(lattices.size());

    for (int i = 0; i < in_spots.size(); i++) {
        spots(i, 0u) = in_spots[i].x;
        spots(i, 1u) = in_spots[i].y;
        spots(i, 2u) = in_spots[i].z;
    }

    for (int i = 0; i < lattices.size(); i++) {
        oCell(i * 3u, 0u) = lattices[i].Vec0().x;
        oCell(i * 3u, 1u) = lattices[i].Vec0().y;
        oCell(i * 3u, 2u) = lattices[i].Vec0().z;

        oCell(i * 3u + 1, 0u) = lattices[i].Vec1().x;
        oCell(i * 3u + 1, 1u) = lattices[i].Vec1().y;
        oCell(i * 3u + 1, 2u) = lattices[i].Vec1().z;

        oCell(i * 3u + 2, 0u) = lattices[i].Vec2().x;
        oCell(i * 3u + 2, 1u) = lattices[i].Vec2().y;
        oCell(i * 3u + 2, 2u) = lattices[i].Vec2().z;

        // Bootstrap score
        scores(i) = 0.2;
    }

    fast_feedback::refine::config_ifssr<float> cifssr{
            .min_spots = viable_cell_min_spots
    };

    fast_feedback::refine::indexer_ifssr<float>::refine(spots.topRows(nspots), oCell, scores, cifssr);

    // Select cell that explains most spots
    int64_t max_indexed_spot_count = 0;
    float max_score = -1;
    int64_t id = -1;

    for (int i = 0; i < lattices.size(); i++) {
        Eigen::Vector3f row_norms = oCell.block(3u * i, 0u, 3u, 3u).rowwise().norm();

        if (row_norms.minCoeff() < min_length_A || row_norms.maxCoeff() > max_length_A)
            scores(i) += 1000;

        auto cell = oCell.block(3u * i, 0u, 3u, 3u);
        M3x resid = spots.topRows(nspots) * cell.transpose();
        const M3x miller = round(resid.array());
        resid -= miller;

        int64_t indexed_spot_count = (resid.rowwise().norm().array() < indexing_tolerance).count();
        if (indexed_spot_count > max_indexed_spot_count) {
            max_indexed_spot_count = indexed_spot_count;
            max_score = scores(i);
            id = i;
        } if (indexed_spot_count == max_indexed_spot_count && scores(i) < max_score) {
          max_score = scores(i);
          id = i;
        }
    }

    if (id == -1)
        return {};

    auto cell = oCell.block(3u * id, 0u, 3u, 3u);

    fast_feedback::refine::make_right_handed(cell);

    return { CrystalLattice(
            Coord(cell(0,0), cell(0,1), cell(0,2)),
            Coord(cell(1,0), cell(1,1), cell(1,2)),
            Coord(cell(2,0), cell(2,1), cell(2,2))
    )};
}