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

#include <cstdint>
#include <vector>
#include "AnalyzeIndexing.h"

#include "FitProfileRadius.h"

namespace {
    inline bool ok(float x) {
        if (!std::isfinite(x))
            return false;
        if (x < 0.0)
            return false;
        return true;
    }

    inline float deg_to_rad(float deg) {
        return deg * (static_cast<float>(M_PI) / 180.0f);
    }

    inline float rad_to_deg(float rad) {
        return rad * (180.0f / static_cast<float>(M_PI));
    }

    // Wrap to [-180, 180] (useful for residuals)
    inline float wrap_deg_pm180(float deg) {
        while (deg > 180.0f) deg -= 360.0f;
        while (deg < -180.0f) deg += 360.0f;
        return deg;
    }

    // Solve A cos(phi) + B sin(phi) + D = 0, return solutions in [phi0, phi1] (radians)
    inline int solve_trig(float A, float B, float D,
                          float phi0, float phi1,
                          float out_phi[2]) {
        const float R = std::sqrt(A * A + B * B);
        if (!(R > 0.0f))
            return 0;

        const float rhs = -D / R;
        if (rhs < -1.0f || rhs > 1.0f)
            return 0;

        const float phi_ref = std::atan2(B, A);
        const float delta = std::acos(rhs);

        float s1 = phi_ref + delta;
        float s2 = phi_ref - delta;

        const float two_pi = 2.0f * static_cast<float>(M_PI);
        auto shift_near = [&](float x) {
            while (x < phi0 - two_pi) x += two_pi;
            while (x > phi1 + two_pi) x -= two_pi;
            return x;
        };

        s1 = shift_near(s1);
        s2 = shift_near(s2);

        int n = 0;
        if (s1 >= phi0 && s1 <= phi1) out_phi[n++] = s1;
        if (s2 >= phi0 && s2 <= phi1) {
            if (n == 0 || std::fabs(s2 - out_phi[0]) > 1e-6f) out_phi[n++] = s2;
        }
        return n;
    }

    // Find predicted phi (deg) for given g0 around phi_obs (deg) within +/- half_window_deg.
    // Returns nullopt if no solution in the local window.
    inline std::optional<float> predict_phi_deg_local(const Coord &g0,
                                                      const Coord &S0,
                                                      const Coord &w_unit,
                                                      float phi_obs_deg,
                                                      float half_window_deg) {
        const float phi0 = deg_to_rad(phi_obs_deg - half_window_deg);
        const float phi1 = deg_to_rad(phi_obs_deg + half_window_deg);

        // Decompose g0 into parallel/perp to w
        const float g_par_s = g0 * w_unit;
        const Coord g_par = w_unit * g_par_s;
        const Coord g_perp = g0 - g_par;

        const float g_perp2 = g_perp * g_perp;
        if (g_perp2 < 1e-12f)
            return std::nullopt;

        const float k2 = (S0 * S0); // |S0|^2 = (1/lambda)^2

        // Equation: |S0 + g(phi)|^2 = |S0|^2
        const Coord p = S0 + g_par;
        const Coord w_x_gperp = w_unit % g_perp;

        const float A = 2.0f * (p * g_perp);
        const float B = 2.0f * (p * w_x_gperp);
        const float D = (p * p) + g_perp2 - k2;

        float sols[2]{};
        const int nsol = solve_trig(A, B, D, phi0, phi1, sols);
        if (nsol == 0)
            return std::nullopt;

        // Pick the solution closest to phi_obs
        const float phi_obs = deg_to_rad(phi_obs_deg);
        float best_phi = sols[0];
        float best_err = std::fabs(sols[0] - phi_obs);

        if (nsol == 2) {
            const float err2 = std::fabs(sols[1] - phi_obs);
            if (err2 < best_err) {
                best_err = err2;
                best_phi = sols[1];
            }
        }

        return rad_to_deg(best_phi);
    }
} // namespace

bool AnalyzeIndexing(DataMessage &message,
                     const DiffractionExperiment &experiment,
                     const CrystalLattice &latt,
                     const std::optional<GoniometerAxis> &rotation_axis) {
    size_t nspots = message.spots.size();

    uint64_t indexed_spot_count = 0;
    std::vector<uint8_t> indexed_spots(nspots);

    // Check spots
    const Coord a = latt.Vec0();
    const Coord b = latt.Vec1();
    const Coord c = latt.Vec2();

    const Coord astar = latt.Astar();
    const Coord bstar = latt.Bstar();
    const Coord cstar = latt.Cstar();

    const auto geom = experiment.GetDiffractionGeometry();
    const auto indexing_tolerance = experiment.GetIndexingSettings().GetTolerance();
    const auto viable_cell_min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots();

    // identify indexed spots
    for (int i = 0; i < message.spots.size(); i++) {
        auto recip = message.spots[i].ReciprocalCoord(geom);

        float h_fp = recip * a;
        float k_fp = recip * b;
        float l_fp = recip * c;

        float h_frac = h_fp - std::round(h_fp);
        float k_frac = k_fp - std::round(k_fp);
        float l_frac = l_fp - std::round(l_fp);

        float norm_sq = h_frac * h_frac + k_frac * k_frac + l_frac * l_frac;

        Coord recip_pred = std::round(h_fp) * astar + std::round(k_fp) * bstar + std::round(l_fp) * cstar;

        // See indexing_peak_check() in peaks.c in CrystFEL
        if (norm_sq < indexing_tolerance * indexing_tolerance) {
            indexed_spot_count++;
            indexed_spots[i] = 1;
            message.spots[i].dist_ewald_sphere = geom.DistFromEwaldSphere(recip_pred);
            message.spots[i].h = std::lround(h_fp);
            message.spots[i].k = std::lround(k_fp);
            message.spots[i].l = std::lround(l_fp);
        }
    }

    auto spot_count_threshold = std::max<int64_t>(viable_cell_min_spots, std::lround(min_percentage_spots * nspots));

    if (indexed_spot_count >= spot_count_threshold) {
        auto uc = latt.GetUnitCell();
        if (!ok(uc.a) || !ok(uc.b) || !ok(uc.c) || !ok(uc.alpha) || !ok(uc.beta) || !ok(uc.gamma))
            return {};

        message.indexing_result = true;
        assert(indexed_spots.size() == message.spots.size());
        for (int i = 0; i < message.spots.size(); i++)
            message.spots[i].indexed = indexed_spots[i];

        message.profile_radius = FitProfileRadius(message.spots);
        message.spot_count_indexed = indexed_spot_count;
        message.indexing_lattice = latt;
        message.indexing_unit_cell = latt.GetUnitCell();

        message.mosaicity_deg = std::nullopt;
        if (rotation_axis.has_value()) {
            const auto gon_opt = experiment.GetGoniometer();
            if (gon_opt.has_value()) {
                const auto &gon = *gon_opt;

                const Coord w = rotation_axis->GetAxis().Normalize();
                const Coord S0 = geom.GetScatteringVector();
                const float wedge_deg = gon.GetWedge_deg();
                const float start_deg = gon.GetStart_deg();
                const float inc_deg = gon.GetIncrement_deg();

                double sum_sq = 0.0;
                int count = 0;

                for (const auto &s: message.spots) {
                    if (!s.indexed)
                        continue;

                    // Observed angle: use frame center
                    const float image_center = static_cast<float>(s.image) + 0.5f;
                    const float phi_obs_deg = start_deg + inc_deg * image_center;

                    // g0 at phi=0 assumption
                    const Coord g0 = astar * static_cast<float>(s.h)
                                     + bstar * static_cast<float>(s.k)
                                     + cstar * static_cast<float>(s.l);

                    // Local solve window: +/- 1 wedge (easy/robust first try)
                    const auto phi_pred_deg_opt = predict_phi_deg_local(g0, S0, w, phi_obs_deg, wedge_deg);
                    if (!phi_pred_deg_opt.has_value())
                        continue;

                    float dphi = wrap_deg_pm180(phi_obs_deg - phi_pred_deg_opt.value());
                    sum_sq += static_cast<double>(dphi) * static_cast<double>(dphi);
                    count++;
                }

                if (count > 0) {
                    message.mosaicity_deg = static_cast<float>(std::sqrt(sum_sq / static_cast<double>(count)));
                }
            }
        }

        return true;
    }

    message.indexing_result = false;
    return false;
}