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

#include "PostIndexingRefinement.h"

namespace {
    struct config_ifssr final {
        float threshold_contraction = .8; // contract error threshold by this value in every iteration
        float max_distance = .00075; // max distance to reciprocal spots for inliers
        unsigned min_spots = 8; // minimum number of spots to fit against
        unsigned max_iter = 32; // max number of iterations
    };

    static std::pair<float, float> score_parts(float score) noexcept {
        float nsp = -std::floor(score);
        float s = score + nsp;
        return std::make_pair(nsp - 1, s);
    }

    struct RefinedCandidate {
        Eigen::Matrix3f cell;
        float score;
        float volume;
        int64_t indexed_spot_count;
        std::vector<uint8_t> indexed_mask;
    };

    static inline Eigen::MatrixX3<float> CalculateResiduals(
        const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
        const Eigen::Matrix3f &cell) {
        Eigen::MatrixX3<float> miller = (spots * cell).array().round().matrix();
        Eigen::MatrixX3<float> resid = miller * cell.inverse();
        resid -= spots;
        return resid;
    }

    static inline std::vector<uint8_t> ComputeIndexedMask(
        const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
        const Eigen::Matrix3f &cell,
        float indexing_tolerance,
        int64_t &indexed_spot_count) {
        const float indexing_tolerance_sq = indexing_tolerance * indexing_tolerance;
        const Eigen::MatrixX3<float> resid = CalculateResiduals(spots, cell);

        std::vector<uint8_t> mask(spots.rows(), 0);
        indexed_spot_count = 0;

        for (int i = 0; i < spots.rows(); ++i) {
            if (resid.row(i).squaredNorm() < indexing_tolerance_sq) {
                mask[i] = 1;
                indexed_spot_count++;
            }
        }

        return mask;
    }

    template<typename MatX3, typename VecX>
    static void RefineCandidateCells(const Eigen::Ref<const Eigen::MatrixX3<float>> &spots,
                                     Eigen::DenseBase<MatX3> &cells,
                                     Eigen::DenseBase<VecX> &scores,
                                     const config_ifssr &cifssr,
                                     unsigned block = 0, unsigned nblocks = 1) {
        using namespace Eigen;
        using Mx3 = MatrixX3<float>;
        using M3 = Matrix3<float>;

        const unsigned nspots = spots.rows();
        const unsigned ncells = scores.rows();
        VectorX<bool> below{nspots};
        MatrixX3<bool> sel{nspots, 3u};
        Mx3 resid{nspots, 3u};
        Mx3 miller{nspots, 3u};
        M3 cell;

        const unsigned blocksize = (ncells + nblocks - 1u) / nblocks;
        const unsigned startcell = block * blocksize;
        const unsigned endcell = std::min(startcell + blocksize, ncells);

        for (unsigned j = startcell; j < endcell; j++) {
            if (nspots < cifssr.min_spots) {
                scores(j) = float{1.};
                continue;
            }

            cell = cells.block(3u * j, 0u, 3u, 3u).transpose(); // cell: col vectors
            const float scale = cell.colwise().norm().minCoeff();
            float threshold = score_parts(scores[j]).second / scale;

            for (unsigned niter = 1; niter < cifssr.max_iter && threshold > cifssr.max_distance; niter++) {
                miller = (spots * cell).array().round().matrix();
                resid = miller * cell.inverse();
                resid -= spots;

                below = (resid.rowwise().norm().array() < threshold);
                if (below.count() < cifssr.min_spots)
                    break;

                threshold *= cifssr.threshold_contraction;
                sel.colwise() = below;
                HouseholderQR<Mx3> qr{sel.select(spots, .0f)};
                cell = qr.solve(sel.select(miller, .0f));
            }

            resid = CalculateResiduals(spots, cell);

            ArrayX<float> dist = resid.rowwise().norm();
            auto nth = std::begin(dist) + (cifssr.min_spots - 1);
            std::nth_element(std::begin(dist), nth, std::end(dist));
            scores(j) = *nth;

            cells.block(3u * j, 0u, 3u, 3u) = cell.transpose();
        }
    }
}

std::vector<CrystalLattice> Refine(const std::vector<Coord> &in_spots,
                                   size_t nspots,
                                   Eigen::MatrixX3<float> &oCell,
                                   Eigen::VectorX<float> &scores,
                                   RefineParameters &p) {
    std::vector<CrystalLattice> ret;

    Eigen::MatrixX3<float> spots(in_spots.size(), 3u);

    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;
    }

    config_ifssr cifssr{
        .min_spots = static_cast<uint32_t>(p.viable_cell_min_spots)
    };

    RefineCandidateCells(spots.topRows(nspots), oCell, scores, cifssr);

    std::vector<RefinedCandidate> candidates;

    for (int i = 0; i < scores.size(); i++) {
        auto cell_block = oCell.block(3u * i, 0u, 3u, 3u);
        Eigen::Matrix3f cell = cell_block.transpose();
        Eigen::Vector3f row_norms = cell_block.rowwise().norm();

        if (p.reference_unit_cell) {
            std::array<float, 3> obs = {row_norms(0), row_norms(1), row_norms(2)};
            std::array<float, 3> ref = {
                static_cast<float>(p.reference_unit_cell->a),
                static_cast<float>(p.reference_unit_cell->b),
                static_cast<float>(p.reference_unit_cell->c)
            };
            std::sort(obs.begin(), obs.end());
            std::sort(ref.begin(), ref.end());

            bool lengths_ok = true;
            for (int k = 0; k < 3; ++k) {
                const float denom = std::max(ref[k], REFINE_MIN_REFERENCE_LENGTH_EPSILON);
                const float rel_dev = std::abs(obs[k] - ref[k]) / denom;
                if (rel_dev > p.dist_tolerance_vs_reference) {
                    lengths_ok = false;
                    break;
                }
            }
            if (!lengths_ok)
                continue;
        } else {
            if (row_norms.minCoeff() < p.min_length_A || row_norms.maxCoeff() > p.max_length_A)
                continue;

            float alpha = std::acos(cell_block.row(1).normalized().dot(cell_block.row(2).normalized())) * 180.0f / M_PI;
            float beta = std::acos(cell_block.row(0).normalized().dot(cell_block.row(2).normalized())) * 180.0f / M_PI;
            float gamma = std::acos(cell_block.row(0).normalized().dot(cell_block.row(1).normalized())) * 180.0f / M_PI;

            if (alpha < p.min_angle_deg || alpha > p.max_angle_deg ||
                beta < p.min_angle_deg || beta > p.max_angle_deg ||
                gamma < p.min_angle_deg || gamma > p.max_angle_deg)
                continue;
        }

        int64_t indexed_spot_count = 0;
        auto indexed_mask = ComputeIndexedMask(spots.topRows(nspots), cell, p.indexing_tolerance, indexed_spot_count);

        if (indexed_spot_count < p.viable_cell_min_spots)
            continue;

        candidates.emplace_back(RefinedCandidate{
            .cell = cell,
            .score = scores(i),
            .volume = std::abs(cell.determinant()),
            .indexed_spot_count = indexed_spot_count,
            .indexed_mask = std::move(indexed_mask)
        });
    }

    std::sort(candidates.begin(), candidates.end(),
              [](const RefinedCandidate &a, const RefinedCandidate &b) {
                  const auto max_spots = std::max(a.indexed_spot_count, b.indexed_spot_count);
                  const auto min_spots = std::min(a.indexed_spot_count, b.indexed_spot_count);
                  const bool spot_counts_close = (max_spots > 0)
                                                 && (static_cast<float>(min_spots) / static_cast<float>(max_spots)
                                                     >= REFINE_CANDIDATE_SPOT_COUNT_RATIO_THRESHOLD);

                  if (!spot_counts_close)
                      return a.indexed_spot_count > b.indexed_spot_count;

                  const float max_volume = std::max(a.volume, b.volume);
                  const float min_volume = std::max(std::min(a.volume, b.volume), REFINE_MIN_VOLUME_EPSILON);
                  const bool volume_differs = (max_volume / min_volume) > REFINE_CANDIDATE_VOLUME_RATIO_THRESHOLD;

                  if (volume_differs)
                      return a.volume < b.volume;

                  if (a.score != b.score)
                      return a.score < b.score;

                  return a.indexed_spot_count > b.indexed_spot_count;
              });

    std::vector<RefinedCandidate> accepted;

    for (const auto &candidate: candidates) {
        bool too_similar = false;

        for (const auto &selected: accepted) {
            int64_t overlap = 0;
            for (size_t i = 0; i < candidate.indexed_mask.size(); ++i) {
                if (candidate.indexed_mask[i] && selected.indexed_mask[i])
                    overlap++;
            }

            const int64_t max_set_size = std::max(candidate.indexed_spot_count, selected.indexed_spot_count);
            if (overlap > static_cast<int64_t>(REFINE_CANDIDATE_OVERLAP_RATIO_THRESHOLD
                                               * static_cast<float>(max_set_size))) {
                too_similar = true;
                break;
            }
        }

        if (!too_similar)
            accepted.emplace_back(candidate);
    }

    ret.reserve(accepted.size());

    for (auto &candidate: accepted) {
        auto cell = candidate.cell;
        if (cell.determinant() < .0f)
            cell = -cell;

        ret.emplace_back(
            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))
        );
    }

    return ret;
}