// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RotationIndexer.h"
#include <iostream>
#include "geom_refinement/XtalOptimizer.h"
#include "lattice_search/LatticeSearch.h"

RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer)
    : experiment(x),
      index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
      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());

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

        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]);
                // Run refinement

                DiffractionExperiment experiment_copy(experiment);
                XtalOptimizerData data{
                    .geom = experiment_copy.GetDiffractionGeometry(),
                    .latt = sym_result.conventional,
                    .crystal_system = sym_result.system,
                    .min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
                    .refine_beam_center = true,
                    .refine_distance_mm =  false,
                    .axis = axis_
                };

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

                if (XtalOptimizer(data, v_)) {
                    indexed_lattice = data.latt;
                    std::cout << sym_result.niggli_class << std::endl;
                    auto uc = indexed_lattice->GetUnitCell();
                    std::cout << uc.a << " " << uc.b << " " << uc.c << " " << uc.alpha << " " << uc.beta << " " << uc.gamma
                    << std::endl;
                    std::cout << indexed_lattice->Vec0().x << " " << indexed_lattice->Vec0().y << " " << indexed_lattice->Vec0().z << std::endl;
                    std::cout << indexed_lattice->Vec1().x << " " << indexed_lattice->Vec1().y << " " << indexed_lattice->Vec1().z << std::endl;
                    std::cout << indexed_lattice->Vec2().x << " " << indexed_lattice->Vec2().y << " " << indexed_lattice->Vec2().z << std::endl;

                }
            }
        }
    }

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