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

#include "AssignSpotsToRings.h"

#include <vector>
#include <cmath>
#include <tuple>
#include <algorithm>
#include <iostream>

#include "../../common/CrystalLattice.h"

FindCircleCenterResult FindCircleCenter(const std::vector<SpotToSave> &v, int64_t width, int64_t height, int64_t max_spots) {
    if ((width <= 0) || (height <= 0) || (max_spots <= 0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid image size");

    std::vector<int64_t> vote(width * height, 0);

    // Limit to only first 250 spots, given the algorithm is N^3
    auto task_size = std::min<size_t>(v.size(), 250);

    for (int i = 0; i < task_size; i++) {
        for (int j = i+1; j < task_size; j++) {
            for (int k = j+1; k < task_size; k++) {
                // Calculation via determinants
                float a = v[i].x * (v[j].y - v[k].y) - v[i].y * (v[j].x - v[k].x) + v[j].x * v[k].y - v[k].x * v[j].y;
                if (std::abs(a) < 1e-10)
                    continue; // Points are collinear

                float x1_sq = v[i].x*v[i].x + v[i].y*v[i].y;
                float x2_sq = v[j].x*v[j].x + v[j].y*v[j].y;
                float x3_sq = v[k].x*v[k].x + v[k].y*v[k].y;

                float bx = x1_sq * (v[j].y - v[k].y) + x2_sq * (v[k].y - v[i].y) + x3_sq * (v[i].y - v[j].y);
                float by = x1_sq * (v[k].x - v[j].x) + x2_sq * (v[i].x - v[k].x) + x3_sq * (v[j].x - v[i].x);

                int64_t cx = std::lround(bx / (2.0f * a));
                int64_t cy = std::lround(by / (2.0f * a));

                if ((cx >= 0) && (cx < width) && (cy >= 0) && (cy < height))
                    vote[cx + cy * width]++;
            }
        }
    }

    int64_t total_votes = 0;
    int64_t max_votes = 0;
    int64_t cx = 0;
    int64_t cy = 0;
    for (int64_t y = 0; y < height; y++) {
        for (int64_t x = 0; x < width; x++) {
            total_votes += vote[x + y * width];
            if (vote[x + y * width] > max_votes) {
                max_votes = vote[x + y * width];
                cx = x;
                cy = y;
            }
        }
    }

    return {
            .total_votes = total_votes,
            .votes_for_beam_center = max_votes,
            .x = static_cast<float>(cx),
            .y = static_cast<float>(cy)
    };
}

// Very simple 1D DBSCAN on radii (works for rings)
std::vector<std::vector<int>> ClusterSpotsIntoRings(const std::vector<float>& r, float eps, int minPts) {
    size_t n = r.size();

    std::vector<int> labels(n, -1); // -1 = unvisited, -2 = noise
    int cluster_id = 0;

    for (int i=0; i<n; i++) {
        if (labels[i] != -1) continue; // already visited

        // find neighbors within eps in radius
        std::vector<int> neighbors;
        for (int j=0; j<n; j++) {
            if (std::fabs(r[i] - r[j]) <= eps) neighbors.push_back(j);
        }

        if ((int)neighbors.size() < minPts) {
            labels[i] = -2; // noise
            continue;
        }

        // start new cluster
        labels[i] = cluster_id;
        std::vector<int> seeds = neighbors;
        for (size_t k=0; k<seeds.size(); k++) {
            int j = seeds[k];
            if (labels[j] == -2) labels[j] = cluster_id;
            if (labels[j] != -1) continue;
            labels[j] = cluster_id;

            // expand cluster
            std::vector<int> nbrs2;
            for (int m=0; m<n; m++) {
                if (std::fabs(r[j] - r[m]) <= eps) nbrs2.push_back(m);
            }
            if ((int)nbrs2.size() >= minPts) {
                seeds.insert(seeds.end(), nbrs2.begin(), nbrs2.end());
            }
        }
        cluster_id++;
    }

    // Collect results
    std::vector<std::vector<int>> clusters(cluster_id);
    for (int i=0; i<n; i++) {
        if (labels[i] >= 0)
            clusters[labels[i]].push_back(i);
    }
    return clusters;
}

float median(std::vector<float> v) {
    if (v.empty()) return std::numeric_limits<float>::quiet_NaN();
    size_t n = v.size();
    std::nth_element(v.begin(), v.begin()+n/2, v.end());
    float m = v[n/2];
    if (n % 2 == 0) {
        auto it = std::max_element(v.begin(), v.begin()+n/2);
        m = 0.5f*(m + *it);
    }
    return m;
}

std::vector<RingClusters> AnalyzeClusters(const std::vector<float>& r, const std::vector<std::vector<int>> &clusters) {
    std::vector<RingClusters> ret;

    for (const auto & cluster : clusters) {
        std::vector<float> cluster_r;
        for (const auto &idx : cluster)
            cluster_r.push_back(r[idx]);

        if (cluster_r.size() < 2) continue;
        float m = median(cluster_r);
        ret.push_back({cluster, m, -1});
    }

    // sort by observed radius
    if (!ret.empty())
        std::sort(ret.begin(), ret.end(), [](const RingClusters& a, const RingClusters& b)
        { return a.R_obs < b.R_obs; });
    return ret;
}

std::vector<float> CalculateXtalRings(const UnitCell &cell, int hkl_max) {
    CrystalLattice latt(cell);

    Coord Astar = latt.Astar();
    Coord Bstar = latt.Bstar();
    Coord Cstar = latt.Cstar();

    std::vector<float> u;
    for (int h = 0; h <= hkl_max; h++) {
        for (int k = 0; k <= hkl_max; k++) {
            for (int l = 0; l <= hkl_max; l++) {
                if (h == 0 && k == 0 && l == 0) continue;
                auto p = Astar * h + Bstar * k + Cstar * l;
                float Q = 2.0f * M_PI * p.Length();
                u.push_back(Q);
            }
        }
    }
    std::sort(u.begin(), u.end());

    // Deduplicate (since e.g. (100), (010), (001) all give sqrt(1))
    u.erase(std::unique(u.begin(), u.end(),
                        [](float a, float b){ return std::fabs(a-b) < 1e-6; }),
            u.end());

    return u;
}


std::vector<float> CalculateCubicXtalRings(float a, int hkl_max) {
    return CalculateXtalRings(UnitCell(a,a,a,90,90,90), hkl_max);
}

float GuessDetectorDistance(const DiffractionGeometry& geom, float ring_radius_pxl, float d_A) {
    float sin_theta = geom.GetWavelength_A() / (2 * d_A);
    if (sin_theta < 0 || sin_theta > 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Geometry makes no sense");
    float theta = asinf(sin_theta);
    float radius_mm = ring_radius_pxl * geom.GetPixelSize_mm();
    float det_dist_mm = radius_mm / tanf(2.0f * theta);
    return det_dist_mm;
}

std::vector<RingClusters> GuessInitialGeometry(DiffractionGeometry &geom,  const std::vector<SpotToSave> &v, float calibrant_a_A) {
    // Reset rotations. The model assumes these are very small in any case!
    geom.PoniRot1_rad(0.0).PoniRot2_rad(0.0).PoniRot3_rad(0.0);

    auto center = FindCircleCenter(v);
    if (center.votes_for_beam_center < 20)
        throw JFJochException(JFJochExceptionCategory::CalibrationError, "Beam center not found");

    geom.BeamX_pxl(center.x).BeamY_pxl(center.y);

    std::vector<float> radii(v.size());
    for (int i = 0; i < v.size(); i++)
        radii[i] = std::hypot(v[i].x - center.x, v[i].y - center.y);

    auto clusters = ClusterSpotsIntoRings(radii);
    if (clusters.empty())
        throw JFJochException(JFJochExceptionCategory::CalibrationError, "Couldn't find spot clusters");

    auto cluster_annot = AnalyzeClusters(radii, clusters);

    float det_distance = GuessDetectorDistance(geom, cluster_annot[0].R_obs, calibrant_a_A);
    geom.DetectorDistance_mm(det_distance);
    return cluster_annot;
}

void GuessGeometry(DiffractionGeometry &geom,  const std::vector<SpotToSave> &v, const UnitCell &calibrant_uc) {
    std::vector<float> ring_q = CalculateXtalRings(calibrant_uc);
    if (ring_q.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Unit cell problem");
    auto cluster_annot = GuessInitialGeometry(geom, v, 2 * M_PI / ring_q[0]);

    std::vector<RingOptimizerInput> optimizer_input;

    int ring_idx = 0;
    int cluster_idx = 0;

    while (cluster_idx < cluster_annot.size() && ring_idx < ring_q.size()) {

        float obs_q = 2 * M_PI / geom.PxlToRes(cluster_annot[cluster_idx].R_obs);
        if (std::fabs(ring_q[ring_idx] - obs_q) < 0.1) {
            std::cout << "Found ring " << ring_idx
                    << " with q_ring = " << ring_q[ring_idx]
                    << " and q_obs = " << obs_q
                    << " diff = " << std::fabs(ring_q[ring_idx] - obs_q) << std::endl;
            for (const auto &spot: cluster_annot[cluster_idx].spots)
                optimizer_input.push_back({v[spot].x, v[spot].y, ring_q[ring_idx]});
            ring_idx++;
            cluster_idx++;
        } else {
            std::cout << "Cannot match " << ring_idx
                    << " with q_ring = " << ring_q[ring_idx]
                    << " and q_obs = " << obs_q
                    << " diff = " << std::fabs(ring_q[ring_idx] - obs_q) << std::endl;
            if (ring_q[ring_idx] < obs_q)
                ring_idx++;
            else
                cluster_idx++;
        }
    }

    RingOptimizer optimizer(geom);
    geom = optimizer.Run(optimizer_input);
}

void OptimizeGeometry(DiffractionGeometry &geom, const std::vector<SpotToSave> &v, const UnitCell &calibrant) {
    std::vector<float> q_ring = CalculateXtalRings(calibrant);

    std::vector<RingOptimizerInput> optimizer_input;

    for (const auto& s: v) {
        float q_obs = 2 * M_PI / geom.PxlToRes(s.x, s.y);
        for (const auto &q : q_ring) {
            if (std::fabs(q - q_obs) < 0.1) {
                optimizer_input.push_back({s.x, s.y, q});
                break;
            }
        }
    }

    RingOptimizer optimizer(geom);
    geom = optimizer.Run(optimizer_input);
}