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

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

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

#include "RotationIndexer.h"

#include "lattice_search/LatticeSearch.h"

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      image_spot0(x.GetImageNum(), 0),
      image_nspots(x.GetImageNum(), 0),
      axis_(x.GetGoniometer()),
      geom_(x.GetDiffractionGeometry()),
      indexer_(indexer) {

    if (axis_) {
        float angle_norm_deg = std::fabs(axis_->GetIncrement_deg());
        if (angle_norm_deg < 1e-6) {
            // Guard against rotation close to zero
            axis_ = std::nullopt;
        } else {
            if (x.GetImageNum() < min_images_for_indexing) {
                // For short measurements - only indexing at the end
                first_image_to_try_indexing = INT64_MAX;
                image_stride = 1;
            } else {
                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 RotationIndexer::SetLattice(const CrystalLattice &lattice) {
    std::unique_lock ul(m);
    indexed_lattice = lattice;
}

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

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

    const auto rot = axis_->GetTransformation(image);

    if (!indexed_lattice && image >= last_accumulated_image + image_stride) {
        v_.reserve(v_.size() + spots.size());
        coords_.reserve(coords_.size() + spots.size());

        image_spot0[image] = v_.size() - 1;
        int i = 0;
        for (const auto &s: spots) {
            if (index_ice_rings || !s.ice_ring) {
                v_.emplace_back(s);
                coords_.emplace_back(rot * s.ReciprocalCoord(geom_));
                i++;
            }
        }
        image_nspots[image] = i;

        accumulated_images++;
        last_accumulated_image = image;

        if (accumulated_images >= min_images_for_indexing && image >= first_image_to_try_indexing) {
            // Index
            auto indexer_result = indexer_.Run(experiment, coords_).get();
            if (!indexer_result.lattice.empty()) {
                // Find lattice type
                auto sym_result = LatticeSearch(indexer_result.lattice[0]);
                indexed_lattice = sym_result.conventional;
                // Run refinement
                //...
            }
        }
    }

    if (indexed_lattice)
        return indexed_lattice->Multiply(rot);
    return {};
}