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

#include "SpotAnalyze.h"

#include "geom_refinement/XtalOptimizer.h"
#include "spot_finding/SpotUtils.h"
#include "spot_finding/StrongPixelSet.h"
#include "bragg_integration/BraggIntegrate2D.h"
#include "indexing/AnalyzeIndexing.h"
#include "bragg_integration/CalcISigma.h"
#include "lattice_search/LatticeSearch.h"

void SpotAnalyze(const DiffractionExperiment &experiment,
                 const SpotFindingSettings &spot_finding_settings,
                 const std::vector<DiffractionSpot> &spots,
                 const CompressedImage &image,
                 BraggPrediction &prediction,
                 IndexerThreadPool *indexer,
                 DataMessage &output) {
    std::vector<SpotToSave> spots_out;

    auto geom = experiment.GetDiffractionGeometry();
    for (const auto &spot: spots)
        spots_out.push_back(spot.Export(geom));

    if (spot_finding_settings.high_res_gap_Q_recipA.has_value())
        FilterSpuriousHighResolutionSpots(spots_out, spot_finding_settings.high_res_gap_Q_recipA.value());

    if (experiment.GetDatasetSettings().IsDetectIceRings() && spot_finding_settings.ice_ring_width_Q_recipA > 0.0f)
        MarkIceRings(spots_out, spot_finding_settings.ice_ring_width_Q_recipA);

    CountSpots(output, spots_out, spot_finding_settings.cutoff_spot_count_low_res);

    GenerateSpotPlot(output, spot_finding_settings.high_resolution_limit);

    output.resolution_estimate = GetResolution(spots_out);

    FilterSpotsByCount(spots_out, experiment.GetMaxSpotCount());

    output.spots = spots_out;

    if ((indexer != nullptr) && spot_finding_settings.indexing) {
        auto latt_f = indexer->Run(experiment, output);
        auto latt = latt_f.get();

        if (!latt)
            output.indexing_result = false;
        else {
            auto uc = latt->GetUnitCell();

            bool beam_center_updated = false;

            auto sg = experiment.GetGemmiSpaceGroup();
            LatticeMessage symmetry{
                .centering = 'P',
                .niggli_class = 0,
                .crystal_system = gemmi::CrystalSystem::Triclinic
            };

            // 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.value());
                symmetry = LatticeMessage{
                    .centering = sym_result.centering,
                    .niggli_class = sym_result.niggli_class,
                    .crystal_system = sym_result.system,
                    .primitive = sym_result.primitive_reduced
                };
                output.lattice_type = symmetry;
                latt = sym_result.conventional;
            }

            DiffractionExperiment experiment_copy(experiment);
            XtalOptimizerData data{
                .geom = experiment_copy.GetDiffractionGeometry(),
                .latt = latt.value(),
                .crystal_system = symmetry.crystal_system,
                .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
            };

            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, output.spots)) {
                        experiment_copy.BeamX_pxl(data.geom.GetBeamX_pxl()).BeamY_pxl(data.geom.GetBeamY_pxl());
                        beam_center_updated = true;
                        latt = data.latt;
                    }
                    break;
            }

            if (AnalyzeIndexing(output, experiment_copy, latt.value())) {
                float ewald_dist_cutoff = 0.001f;

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

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

                if (spot_finding_settings.quick_integration) {
                    auto res = BraggIntegrate2D(experiment_copy, image, latt.value(),
                                                prediction, ewald_dist_cutoff, output.number, symmetry.centering);

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

                    CalcISigma(output);
                    CalcWilsonBFactor(output);
                }
            }
        }
    }
}