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

#include "IndexAndRefine.h"

#include "bragg_integration/BraggIntegrate2D.h"
#include "bragg_integration/CalcISigma.h"
#include "geom_refinement/XtalOptimizer.h"
#include "indexing/AnalyzeIndexing.h"
#include "indexing/FFTIndexer.h"
#include "lattice_search/LatticeSearch.h"

IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer)
    : min_accum_angle_deg(x.GetIndexingSettings().GetRotationIndexingMinAngularRange_deg()),
      stride_angle_deg(x.GetIndexingSettings().GetRotationIndexingAngularStride_deg()),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      experiment(x),
      geom_(x.GetDiffractionGeometry()),
      updated_geom_(geom_),
      indexer_(indexer) {
    auto axis = x.GetGoniometer();

    if (indexer && axis && x.GetIndexingSettings().GetRotationIndexing()) {
        float angle_norm_deg = std::fabs(axis->GetIncrement_deg());
        if (angle_norm_deg > 1e-6) {
            axis_ = axis;
            if (x.GetImageNum() > min_images_for_indexing) {
                first_image_to_try_indexing = std::max<int64_t>(min_images_for_indexing,
                                                                min_accum_angle_deg / angle_norm_deg);
                image_stride = std::ceil(stride_angle_deg / angle_norm_deg);
                if (image_stride == 0)
                    image_stride = 1;
            }
        }
    }
}

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

void IndexAndRefine::TryIndexRotationData() {
    if (!indexer_)
        return;

    // 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) {
                // Replace `.intensity` with the actual SpotToSave intensity member
                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 = true,
            .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;
        }
    }
}

void IndexAndRefine::ProcessImage(DataMessage &msg,
                                  const SpotFindingSettings &spot_finding_settings,
                                  const CompressedImage &image,
                                  BraggPrediction &prediction) {
    if (!indexer_)
        return;

    msg.indexing_result = false;
    std::optional<CrystalLattice> lattice_candidate;
    DiffractionExperiment experiment_copy(experiment);

    LatticeMessage symmetry{
        .centering = 'P',
        .niggli_class = 0,
        .crystal_system = gemmi::CrystalSystem::Triclinic
    };

    bool beam_center_updated = false;

    if (!axis_) {
        // Convert input spots to reciprocal space
        std::vector<Coord> recip;
        recip.reserve(msg.spots.size());
        for (const auto &i: msg.spots) {
            if (index_ice_rings || !i.ice_ring)
                recip.push_back(i.ReciprocalCoord(updated_geom_));
        }

        auto indexer_result = indexer_->Run(experiment, recip).get();
        msg.indexing_time_s = indexer_result.indexing_time_s;

        if (!indexer_result.lattice.empty()) {
            auto latt = indexer_result.lattice[0];

            auto sg = experiment.GetGemmiSpaceGroup();

            // If space group provided => always enforce symmetry in refinement
            // If space group not provided => guess symmetry
            if (sg) {
                // If space group provided but no unit cell fixed, it is better to keep triclinic for now
                if (experiment.GetUnitCell()) {
                    symmetry = LatticeMessage{
                        .centering = sg->centring_type(),
                        .niggli_class = 0,
                        .crystal_system = sg->crystal_system()
                    };
                }
            } else {
                auto sym_result = LatticeSearch(latt);
                symmetry = LatticeMessage{
                    .centering = sym_result.centering,
                    .niggli_class = sym_result.niggli_class,
                    .crystal_system = sym_result.system
                };

                latt = sym_result.conventional;
            }

            lattice_candidate = latt;
        }
    } else {
        std::unique_lock ul(m);
        const auto rot = axis_->GetTransformation(msg.number);

        if (msg.number >= last_accumulated_image + image_stride) {
            v_.reserve(v_.size() + msg.spots.size());
            coords_.reserve(coords_.size() + msg.spots.size());

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

            accumulated_images++;
            last_accumulated_image = msg.number;

            if (!indexed_lattice
                && accumulated_images >= min_images_for_indexing
                && msg.number >= first_image_to_try_indexing)
                TryIndexRotationData();
        }

        if (indexed_lattice) {
            lattice_candidate = indexed_lattice->Multiply(rot.transpose());
            symmetry = LatticeMessage{
                .centering = search_result_.centering,
                .niggli_class = search_result_.niggli_class,
                .crystal_system = search_result_.system
            };
            beam_center_updated = true;
        }
    }

    if (lattice_candidate) {
        XtalOptimizerData data{
            .geom = updated_geom_,
            .latt = lattice_candidate.value(),
            .crystal_system = symmetry.crystal_system,
            .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
            .refine_beam_center = true,
            .refine_distance_mm = false
        };

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

        switch (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm()) {
            case GeomRefinementAlgorithmEnum::None:
                break;
            case GeomRefinementAlgorithmEnum::BeamCenter:
                if (XtalOptimizer(data, msg.spots)) {
                    experiment_copy.BeamX_pxl(data.geom.GetBeamX_pxl()).BeamY_pxl(data.geom.GetBeamY_pxl());
                    beam_center_updated = true;
                    lattice_candidate = data.latt;
                }
                break;
        }
        if (XtalOptimizer(data, msg.spots)) {
            if (AnalyzeIndexing(msg, experiment, data.latt)) {
                msg.lattice_type = symmetry;

                float ewald_dist_cutoff = 0.001f;

                if (msg.profile_radius)
                    ewald_dist_cutoff = msg.profile_radius.value() * 2.0f;
                if (experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA())
                    ewald_dist_cutoff = experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA().value()
                                        * 3.0f;

                if (beam_center_updated) {
                    msg.beam_corr_x = data.beam_corr_x;
                    msg.beam_corr_y = data.beam_corr_y;
                }

                if (spot_finding_settings.quick_integration) {
                    auto res = BraggIntegrate2D(experiment_copy, image, *lattice_candidate,
                                                prediction, ewald_dist_cutoff, msg.number, symmetry.centering);

                    constexpr size_t kMaxReflections = 10000;
                    if (res.size() > kMaxReflections) {
                        msg.reflections.assign(res.begin(), res.begin() + kMaxReflections);
                    } else
                        msg.reflections = res;

                    CalcISigma(msg);
                    CalcWilsonBFactor(msg);
                }
            }
        }
    }
}

std::optional<RotationIndexerResult> IndexAndRefine::Finalize(bool retry) {
    if (axis_) {
        std::unique_lock ul(m);

        if (!indexed_lattice || retry)
            TryIndexRotationData();
        if (indexed_lattice)
            return RotationIndexerResult{
                .lattice = indexed_lattice.value(),
                .search_result = search_result_,
                .geom = geom_
            };
    }
    return {};
}