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

#pragma once

#include <vector>
#include <mutex>

#include "../common/DiffractionSpot.h"
#include "../common/DiffractionExperiment.h"
#include "bragg_prediction/BraggPrediction.h"
#include "indexing/IndexerThreadPool.h"
#include "lattice_search/LatticeSearch.h"
#include "scale_merge/Merge.h"
#include "RotationIndexer.h"
#include "RotationParameters.h"
#include "scale_merge/ScaleOnTheFly.h"
#include "scale_merge/ScalingResult.h"

class IndexAndRefine {
    const bool index_ice_rings;
    const DiffractionExperiment& experiment;
    const DiffractionGeometry geom_;

    std::optional<CrystalLattice> indexed_lattice;

    std::optional<GoniometerAxis> axis_;

    IndexerThreadPool *indexer_;
    std::unique_ptr<RotationIndexer> rotation_indexer;

    RotationParameters rotation_parameters;

    struct IndexingOutcome {
        std::optional<CrystalLattice> lattice_candidate;
        DiffractionExperiment experiment;
        LatticeMessage symmetry{
            .centering = 'P',
            .niggli_class = 0,
            .crystal_system = gemmi::CrystalSystem::Triclinic
        };
        bool beam_center_updated = false;

        explicit IndexingOutcome(const DiffractionExperiment& experiment_ref)
            : experiment(experiment_ref) {}
    };

    mutable std::mutex reflections_mutex;
    std::vector<std::vector<Reflection>> reflections;
    std::vector<float> mosaicity;
    std::vector<float> scale_cc;
    std::vector<std::optional<UnitCell> > unit_cells;

    IndexingOutcome DetermineLatticeAndSymmetryRotation(DataMessage &msg);
    IndexingOutcome DetermineLatticeAndSymmetry(DataMessage &msg);
    void RefineGeometryIfNeeded(DataMessage &msg, IndexingOutcome &outcome);
    void QuickPredictAndIntegrate(DataMessage &msg,
                                 const SpotFindingSettings &spot_finding_settings,
                                 const CompressedImage &image,
                                 BraggPrediction &prediction,
                                 const IndexingOutcome &outcome);

    std::unique_ptr<ScaleOnTheFly> scaling_engine;
    void ScaleImage(DataMessage &msg);
public:
    IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer);
    void ProcessImage(DataMessage &msg, const SpotFindingSettings &settings, const CompressedImage &image, BraggPrediction &prediction);
    IndexAndRefine& ReferenceIntensities(std::vector<MergedReflection> &reference);

    ScalingResult ScaleAllImages(const std::vector<MergedReflection> &reference, size_t nthreads = 0);

    std::optional<RotationIndexerResult> Finalize();

    std::optional<UnitCell> GetConsensusUnitCell() const;

    // Not thread safe, need to be run after processing is all done
    const std::vector<std::vector<Reflection>> &GetReflections() const;
    std::vector<std::vector<Reflection>> &GetReflections();
    const std::vector<float> &GetImageCC() const;
    const std::vector<std::optional<UnitCell> > &GetUnitCells() const;


};