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

#include "FFTIndexer.h"
#include <Eigen/Eigen>
#include "EigenRefine.h"

FFTIndexer::FFTIndexer(const IndexingSettings &settings)
    : max_length_A(settings.GetFFT_MaxUnitCell_A()),
      min_length_A(settings.GetFFT_MinUnitCell_A()),
      min_angle_deg(settings.GetFFT_MinAngle_deg()),
      max_angle_deg(settings.GetFFT_MaxAngle_deg()),
      nDirections(settings.GetFFT_NumVectors()),
      result_fft(nDirections) {

    float maxQ = 2.0f * static_cast<float>(M_PI) / settings.GetFFT_HighResolution_A();

    histogram_spacing = 1.0f / (2.0f * max_length_A);
    histogram_size = std::ceil(maxQ / histogram_spacing);

    input_size = histogram_size * nDirections;
    output_size = (histogram_size / 2 + 1) * nDirections;

    if (max_length_A <= settings.GetFFT_HighResolution_A())
        throw std::invalid_argument("Largest unit cell cannot be smaller than resolution");
    if (nDirections <= 0)
        throw std::invalid_argument("FFTWIndexer: number of directions must be > 0");
    if (!(max_length_A > 0.f))
        throw std::invalid_argument("FFTWIndexer: max_length_A must be > 0");
    if (!(settings.GetFFT_HighResolution_A() > 0.f))
        throw std::invalid_argument("FFTWIndexer: high resolution must be > 0");
    if (histogram_size < 1)
        throw std::invalid_argument("FFTWIndexer: histogram_size must be >= 1");
    if (histogram_size > 1000000)
        throw std::invalid_argument("FFTWIndexer: histogram_size too large");

    SetupDirectionVectors();
}

void FFTIndexer::SetupDirectionVectors() {
    direction_vectors.reserve(static_cast<size_t>(nDirections));

    const double phi = (1.0 + std::sqrt(5.0)) / 2.0; // Golden ratio
    const double golden_angle = 2.0 * M_PI / phi;

    for (int i = 0; i < nDirections; i++) {
        // Half-sphere distribution (z in [0,1])
        double z = static_cast<double>(i) / static_cast<double>(nDirections - 1);
        double theta = golden_angle * static_cast<double>(i);
        double radius = std::sqrt(std::max(0.0, 1.0 - z * z));

        double x = radius * std::cos(theta);
        double y = radius * std::sin(theta);

        // Add unit vector to the list
        direction_vectors.emplace_back(static_cast<float>(x),
                                       static_cast<float>(y),
                                       static_cast<float>(z));
    }
}


void FFTIndexer::SetupUnitCell(const std::optional<UnitCell> &cell) {
    reference_unit_cell = cell;
}


void Sort(Coord &A, Coord &B, Coord &C) {
    // A is the smallest, C is the largest

    if (A.Length() > B.Length())
        std::swap(A, B);

    if (B.Length() > C.Length())
        std::swap(B, C);

    if (A.Length() > B.Length())
        std::swap(A, B);
}

std::vector<CrystalLattice> FFTIndexer::ReduceResults(const std::vector<Coord> &results) const {
    if (results.size() < 3)
        return {};

    std::vector<CrystalLattice> candidates;

    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            for (int k = 0; k < 3; k++) {
                if (i + j + k + 2 >= results.size())
                    break;

                Coord A = results[i];
                Coord B = results[(i + j + 1)];
                Coord C = results[(i + j + 1) + k + 1];

                // sort vectors by length for reduction
                Sort(A,B,C);

                // Lattice reduction
                B = B - std::round(B * A / (A * A)) * A;
                C = C - std::round(C * A / (A * A)) * A;
                C = C - std::round(C * B / (B * B)) * B;

                float alpha = angle_deg(B, C);
                float beta = angle_deg(A, C);
                float gamma = angle_deg(A, B);

                // Check if values are OK after lattice reduction
                if (A.Length() >= min_length_A
                    && B.Length() >= min_length_A
                    && C.Length() >= min_length_A
                    && alpha >= min_angle_deg && alpha <= max_angle_deg
                    && beta >= min_angle_deg && beta <= max_angle_deg
                    && gamma >= min_angle_deg && gamma <= max_angle_deg)
                    candidates.emplace_back(A, B, C);
            }
        }
    }

    return candidates;
}

std::vector<Coord> FFTIndexer::FilterFFTResults() const {
    std::multimap<float, FFTResult> fft_result_map;

    // Order vectors by max FFT magnitude and select the 30 strongest ones
    size_t max_vectors = 30;

    for (int i = 0; i < direction_vectors.size(); i++)
        fft_result_map.insert(std::make_pair(result_fft[i].magnitude, result_fft[i]));

    std::vector<FFTResult> fft_result_filtered;
    int count = 0;
    for (auto it = fft_result_map.rbegin();
         it != fft_result_map.rend() && count < max_vectors;
         ++it, ++count) {
        fft_result_filtered.emplace_back(it->second);
         }

    std::vector<Coord> ret;

    // Remove vectors less than 5 deg apart, as most likely these are colinear
    constexpr float COS_5_DEG = 0.996194;
    std::vector<bool> ignore(fft_result_filtered.size(), false);

    for (int i = 0; i < fft_result_filtered.size(); i++) {
        if (ignore[i])
            continue;
        Coord dir_i = direction_vectors.at(fft_result_filtered[i].direction);
        for (int j = i + 1; j < fft_result_filtered.size(); j++) {
            Coord dir_j = direction_vectors.at(fft_result_filtered[j].direction);
            if (std::fabs(dir_i * dir_j) > COS_5_DEG)
                ignore[j] = true;
        }
        ret.push_back(dir_i * fft_result_filtered[i].length);
    }
    return ret;
}


std::vector<CrystalLattice> FFTIndexer::Run(const std::vector<Coord> &coord, size_t nspots) {
    if (nspots > coord.size())
        nspots = coord.size();

    if (nspots <= viable_cell_min_spots)
        return {};

    assert(nspots <= FFT_MAX_SPOTS);
    assert(coord.size() <= FFT_MAX_SPOTS);

    ExecuteFFT(coord, nspots);
    auto f = FilterFFTResults();
    auto r = ReduceResults(f);

    Eigen::MatrixX3<float> oCell(r.size() * 3u, 3u);
    Eigen::VectorX<float> scores(r.size());

    for (int i = 0; i < r.size(); i++) {
        oCell(i * 3u, 0u) = r[i].Vec0().x;
        oCell(i * 3u, 1u) = r[i].Vec0().y;
        oCell(i * 3u, 2u) = r[i].Vec0().z;

        oCell(i * 3u + 1, 0u) = r[i].Vec1().x;
        oCell(i * 3u + 1, 1u) = r[i].Vec1().y;
        oCell(i * 3u + 1, 2u) = r[i].Vec1().z;

        oCell(i * 3u + 2, 0u) = r[i].Vec2().x;
        oCell(i * 3u + 2, 1u) = r[i].Vec2().y;
        oCell(i * 3u + 2, 2u) = r[i].Vec2().z;

        // Bootstrap score
        scores(i) = 0.2;
    }

    RefineParameters parameters{
        .viable_cell_min_spots = viable_cell_min_spots,
        .dist_tolerance_vs_reference = dist_tolerance_vs_reference,
        .reference_unit_cell = reference_unit_cell,
        .min_length_A = min_length_A,
        .max_length_A = max_length_A,
        .min_angle_deg = min_angle_deg,
        .max_angle_deg = max_angle_deg,
        .indexing_tolerance = indexing_tolerance
    };

    return Refine(coord, nspots, oCell, scores, parameters);
}