// 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)
    : index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      experiment(x),
      geom_(x.GetDiffractionGeometry()),
      indexer_(indexer) {
    if (indexer && x.GetGoniometer() && x.GetIndexingSettings().GetRotationIndexing())
        rotation_indexer = std::make_unique<RotationIndexer>(x, *indexer);
}

void IndexAndRefine::SetLattice(const CrystalLattice &lattice) {
    if (rotation_indexer)
        rotation_indexer->SetLattice(lattice);
}

void IndexAndRefine::ProcessImage(DataMessage &msg,
                                  const SpotFindingSettings &spot_finding_settings,
                                  const CompressedImage &image,
                                  BraggPrediction &prediction) {
    if (!indexer_ || !spot_finding_settings.indexing)
        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 (rotation_indexer) {
        auto result = rotation_indexer->ProcessImage(msg.number, msg.spots);
        if (result) {
            lattice_candidate = result->lattice;
            experiment_copy.BeamX_pxl(result->geom.GetBeamX_pxl())
                             .BeamY_pxl(result->geom.GetBeamY_pxl())
                             .DetectorDistance_mm(result->geom.GetDetectorDistance_mm());

            symmetry.centering = result->search_result.centering;
            symmetry.niggli_class = result->search_result.niggli_class;
            symmetry.crystal_system = result->search_result.system;
        }
    } else {
        // 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(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()
                    };
                }
                lattice_candidate = latt;
            } else {
                auto sym_result = LatticeSearch(latt);
                symmetry = LatticeMessage{
                    .centering = sym_result.centering,
                    .niggli_class = sym_result.niggli_class,
                    .crystal_system = sym_result.system
                };

                lattice_candidate = sym_result.conventional;
            }
        }
    }

    if (lattice_candidate) {
        XtalOptimizerData data{
            .geom = experiment_copy.GetDiffractionGeometry(),
            .latt = *lattice_candidate,
            .crystal_system = symmetry.crystal_system,
            .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
            .refine_beam_center = true,
            .refine_distance_mm = false,
            .max_time = 0.04 // 40 ms is max allowed time for the operation
        };

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

        lattice_candidate = data.latt;

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

        if (AnalyzeIndexing(msg, experiment_copy, *lattice_candidate, experiment_copy.GetGoniometer())) {
            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 (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() {
    if (rotation_indexer)
        return rotation_indexer->GetLattice();
    return {};
}