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

#include "../../common/JFJochMath.h"
#include "RotationIndexer.h"

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

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      v_(experiment.GetImageNum()),
      angle_deg_(experiment.GetImageNum()),
      axis_(x.GetGoniometer()),
      geom_(x.GetDiffractionGeometry()),
      updated_geom_(geom_),
      indexer_(indexer) {
}

void RotationIndexer::RunIndexing() {
    std::unique_lock ul(m);
    if (!axis_)
        return;

    std::vector<Coord> coords;
    coords.reserve(max_spots_per_image * v_.size());

    for (int i = 0; i < v_.size(); i++) {
        const float angle_deg = angle_deg_[i].value_or(axis_->GetAngle_deg(i) + axis_->GetWedge_deg() / 2.0f);
        const auto rot = axis_->GetTransformationAngle(angle_deg);

        for (const auto &s: v_[i])
            coords.emplace_back(rot * s.ReciprocalCoord(geom_));
    }
    const auto indexer_result = indexer_.Run(experiment, coords);

    if (!indexer_result.lattice.empty() && indexer_result.lattice[0].CalcVolume() > 1.0) {
        auto sg = experiment.GetGemmiSpaceGroup();
        if (sg) {
            search_result_ = LatticeSearchResult{
                .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
                .conventional = indexer_result.lattice[0], // If lattice provided, it is for now primitive == conventional
                .system = sg->crystal_system(),
                .centering = sg->centring_type(),
                .reindex = gemmi::Mat33(1, 0, 0, 0, 1, 0, 0, 0, 1),
            };
        } else {
            // Find lattice type based on cell
            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,
            .refine_detector_angles = true,
            .refine_rotation_axis = true,
            .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
            .axis = axis_
        };

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

        if (data.crystal_system == gemmi::CrystalSystem::Monoclinic)
            data.latt.ReorderMonoclinic();

        if (XtalOptimizer(data, v_)) {
            indexed_lattice = data.latt;
            updated_geom_ = data.geom;
            axis_ = data.axis;
        }

        if (indexer_result.lattice.size() > 1) {
            auto ml_latt = MultiLatticeSearch(indexer_result.lattice);
            for (auto &l : ml_latt) {
                if (extra_lattices_.size() >= experiment.GetIndexingSettings().GetMaxExtraLattices())
                    break;

                // Ignore lattices oriented by less than 3.0 degree
                if (l.rotation_vector.Length() < 3.0 * PI / 180.0)
                    continue;

                RotMatrix rot(l.rotation_vector.Length(), l.rotation_vector.Normalize());

                XtalOptimizerData data_multi{
                    .geom = experiment_copy.GetDiffractionGeometry(),
                    .latt = data.latt.Multiply(rot),
                    .crystal_system = search_result_.system,
                    .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                    .refine_beam_center = false,
                    .refine_distance_mm =  false,
                    .refine_detector_angles = false,
                    .refine_unit_cell = false,
                    .refine_rotation_axis = false,
                    .index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
                    .axis = axis_
                };

                // Quick refinement: orientation only. Cell size/angles, beam center,
                // detector angles and rotation axis are all kept from the first lattice.
                // XtalOptimizer always refines orientation; everything else is frozen above.
                XtalOptimizer(data_multi, v_);

                extra_lattices_.push_back(data_multi.latt);
            }
        }
    }
}

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

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

    // Guard: `image` is a slot in [0, image count); a bad index (e.g. a global number for a subset
    // run) must not corrupt memory.
    if (image < 0 || image >= static_cast<int64_t>(v_.size()))
        return;

    if (accumulated_spots >= max_spots)
        return;

    if (indexed_lattice)
        return;

    angle_deg_[image] = angle_deg;
    v_[image].reserve(spots.size());

    for (const auto &s: spots) {
        if (index_ice_rings || !s.ice_ring)
            v_[image].emplace_back(s);
    }

    // truncate spots, so we don't get above max_spots (total) and max_spots_per_image (for this image)
    size_t max_spots_limit = std::min(max_spots_per_image, max_spots - accumulated_spots);

    if (v_[image].size() > max_spots_limit) {
        std::ranges::nth_element(v_[image], v_[image].begin() + max_spots_limit,
                                 [](const SpotToSave &a, const SpotToSave &b) {
                                     return a.intensity > b.intensity;
                                 }
        );

        v_[image].resize(max_spots_limit);
    }

    accumulated_spots += v_[image].size();
}

std::optional<RotationIndexerResult> RotationIndexer::GetLattice() const {
    std::unique_lock ul(m);

    if (!indexed_lattice)
        return {};
    return RotationIndexerResult{
        .lattice = indexed_lattice.value(),
        .extra_lattices = extra_lattices_,
        .search_result = search_result_,
        .geom = updated_geom_,
        .axis = axis_,
    };
}

void RotationIndexer::ForceLattice(const CrystalLattice &lattice) {
    indexed_lattice = lattice;
    auto sg_num = experiment.GetSpaceGroupNumber().value_or(1);
    auto sg = gemmi::find_spacegroup_by_number(sg_num);
    if (sg != nullptr) {
        search_result_ = LatticeSearchResult{
            .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
            .conventional = lattice, // If lattice provided, it is for now primitive == conventional
            .system = sg->crystal_system(),
            .centering = sg->centring_type(),
        };
    } else
        search_result_ = LatticeSearchResult{
            .niggli_class = 0, // Since Niggli class was not searched for, we don't know which one
            .conventional = lattice, // If lattice provided, it is for now primitive == conventional
            .system = gemmi::CrystalSystem::Triclinic,
            .centering = 'P',
        };
}