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

#include "RotationIndexer.h"

#include "geom_refinement/XtalOptimizer.h"
#include "indexing/FFTIndexer.h"
#include "lattice_search/LatticeSearch.h"

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      axis_(x.GetGoniometer()),
      geom_(x.GetDiffractionGeometry()),
      updated_geom_(geom_),
      indexer_(indexer) {
    if (axis_) {
        float angle_norm_deg = std::fabs(axis_->GetIncrement_deg());
        if (angle_norm_deg < 1e-6) {
            // Guard against rotation close to zero
            axis_ = std::nullopt;
        } else {
            if (x.GetImageNum() < min_images_for_indexing) {
                // For short measurements - only indexing at the end
                first_image_to_try_indexing = INT64_MAX;
                image_stride = 1;
            } else {
                first_image_to_try_indexing = std::max<int64_t>(min_images_for_indexing,
                    x.GetIndexingSettings().GetRotationIndexingMinAngularRange_deg() / angle_norm_deg);
                image_stride = std::ceil(x.GetIndexingSettings().GetRotationIndexingAngularStride_deg() / angle_norm_deg);
                if (image_stride == 0)
                    image_stride = 1;
            }
        }
    }
}

void RotationIndexer::SetLattice(const CrystalLattice &lattice) {
    std::unique_lock ul(m);
    indexed_lattice = lattice;
}

void RotationIndexer::TryIndex() {
    indexing_tried = true;

    // Index
    std::vector<SpotToSave> v_sel;
    std::vector<Coord> coords_sel;

    if (coords_.size() > max_spots) {

        // Indices into v_ / coords_
        std::vector<std::size_t> idx(coords_.size());
        std::iota(idx.begin(), idx.end(), std::size_t{0});

        // Sort indices by descending intensity
        std::partial_sort(
            idx.begin(),
            idx.begin() + max_spots,
            idx.end(),
            [this](std::size_t a, std::size_t b) {
                return v_[a].intensity > v_[b].intensity;
            }
        );

        v_sel.reserve(max_spots);
        coords_sel.reserve(max_spots);

        for (std::size_t i = 0; i < max_spots; ++i) {
            const std::size_t k = idx[i];
            v_sel.emplace_back(v_[k]);
            coords_sel.emplace_back(coords_[k]);
        }
    } else {
        v_sel = v_;
        coords_sel = coords_;
    }

    auto indexer_result = indexer_.Run(experiment, coords_sel).get();
    if (!indexer_result.lattice.empty()) {
        // Find lattice type
        search_result_ = LatticeSearch(indexer_result.lattice[0]);
        // Run refinement

        DiffractionExperiment experiment_copy(experiment);
        XtalOptimizerData data{
            .geom = experiment_copy.GetDiffractionGeometry(),
            .latt = search_result_.conventional,
            .crystal_system = search_result_.system,
            .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
            .refine_beam_center = true,
            .refine_distance_mm =  false,
            .axis = axis_
        };

        if (data.crystal_system == gemmi::CrystalSystem::Trigonal)
            data.crystal_system = gemmi::CrystalSystem::Hexagonal;

        if (XtalOptimizer(data, v_sel)) {
            indexed_lattice = data.latt;
            updated_geom_ = data.geom;
        }
    }
}

std::optional<RotationIndexerResult> RotationIndexer::ProcessImage(int64_t image, const std::vector<SpotToSave> &spots) {
    std::unique_lock ul(m);

    // For non-rotation just ignore the whole procedure
    if (!axis_)
        return {};

    const auto rot = axis_->GetTransformation(image);

    if (!indexing_tried && image >= last_accumulated_image + image_stride) {
        v_.reserve(v_.size() + spots.size());
        coords_.reserve(coords_.size() + spots.size());

        for (const auto &s: spots) {
            if (index_ice_rings || !s.ice_ring) {
                v_.emplace_back(s);
                coords_.emplace_back(rot * s.ReciprocalCoord(geom_));
            }
        }

        accumulated_images++;
        last_accumulated_image = image;

        if (accumulated_images >= min_images_for_indexing && image >= first_image_to_try_indexing)
            TryIndex();
    }

    return GetLattice();
}

std::optional<RotationIndexerResult> RotationIndexer::GetLattice() {
    if (!indexed_lattice)
        return {};
    return RotationIndexerResult{
        .lattice = indexed_lattice.value(),
        .search_result = search_result_,
        .geom = updated_geom_
    };
}