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

#include <cstdint>
#include <vector>
#include <cmath>
#include <algorithm>

#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;
    }

     // XDS convention: zeta = |m2 · e1| where e1 = (S × S0) / |S × S0|
    // This is the Lorentz factor component related to the rotation axis
    inline float calc_zeta(const Coord& S, const Coord& S0, const Coord& m2) {
        Coord S_cross_S0 = S % S0;
        float len = S_cross_S0.Length();
        if (len < 1e-12f) return 0.0f;
        Coord e1 = S_cross_S0 * (1.0f / len);
        return std::fabs(m2 * e1);
    }

    // XDS R(τ; σM/ζ) function - fraction of observed reflection intensity
    // τ = angular difference between reflection and Bragg maximum (radians)
    // delta_phi = oscillation range (radians)
    // sigma_M = mosaicity (radians)
    // zeta = |m2 · e1| Lorentz factor component
    inline float R_fraction(float tau, float delta_phi, float sigma_M, float zeta) {
        if (zeta < 1e-6f || sigma_M < 1e-9f)
            return 0.0f;

        const float sigma_eff = sigma_M / zeta;
        const float sqrt2_sigma = std::sqrt(2.0f) * sigma_eff;

        if (sqrt2_sigma < 1e-12f)
            return 0.0f;

        const float arg_plus = (tau + delta_phi / 2.0f) / sqrt2_sigma;
        const float arg_minus = (tau - delta_phi / 2.0f) / sqrt2_sigma;

        return 0.5f * (std::erf(arg_plus) - std::erf(arg_minus));
    }

    // Log-likelihood for a given sigma_M value
    // Returns sum of log(R) for all reflections
    inline double log_likelihood(const std::vector<float>& tau_values,
                                 const std::vector<float>& zeta_values,
                                 float delta_phi,
                                 float sigma_M) {
        double ll = 0.0;
        for (size_t i = 0; i < tau_values.size(); ++i) {
            float R = R_fraction(tau_values[i], delta_phi, sigma_M, zeta_values[i]);
            if (R > 1e-30f) {
                ll += std::log(static_cast<double>(R));
            } else {
                ll += -70.0;  // Large penalty for zero probability
            }
        }
        return ll;
    }

    // Golden section search for maximum likelihood sigma_M
    inline float find_sigma_M_mle(const std::vector<float>& tau_values,
                                  const std::vector<float>& zeta_values,
                                  float delta_phi,
                                  float sigma_min_deg = 0.001f,
                                  float sigma_max_deg = 2.0f) {
        const float golden = 0.618033988749895f;

        float a = deg_to_rad(sigma_min_deg);
        float b = deg_to_rad(sigma_max_deg);

        float c = b - golden * (b - a);
        float d = a + golden * (b - a);

        const float tol = 1e-6f;

        while (std::fabs(b - a) > tol) {
            double fc = log_likelihood(tau_values, zeta_values, delta_phi, c);
            double fd = log_likelihood(tau_values, zeta_values, delta_phi, d);

            if (fc > fd) {
                b = d;
                d = c;
                c = b - golden * (b - a);
            } else {
                a = c;
                c = d;
                d = a + golden * (b - a);
            }
        }

        return (a + b) / 2.0f;
    }

    // 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);
    }

    // XDS-style mosaicity calculation using maximum likelihood
    // Following Kabsch (2010) Acta Cryst. D66, 133-144
    std::optional<float> CalcMosaicityXDS(const DiffractionExperiment& experiment,
                                          const std::vector<SpotToSave> &spots,
                                          const Coord &astar, const Coord &bstar, const Coord &cstar) {
        const auto &axis_opt = experiment.GetGoniometer();
        if (!axis_opt.has_value())
            return std::nullopt;

        const GoniometerAxis& axis = *axis_opt;
        const Coord m2 = axis.GetAxis().Normalize();  // XDS notation: m2 is rotation axis
        const Coord S0 = experiment.GetScatteringVector();
        const float delta_phi_rad = deg_to_rad(axis.GetWedge_deg());

        std::vector<float> tau_values;
        std::vector<float> zeta_values;
        tau_values.reserve(spots.size());
        zeta_values.reserve(spots.size());

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

            const Coord pstar = astar * static_cast<float>(s.h)
                              + bstar * static_cast<float>(s.k)
                              + cstar * static_cast<float>(s.l);

            // Find predicted phi angle
            const auto phi_pred_opt = predict_phi_deg_local(pstar, S0, m2, 0.0f, axis.GetWedge_deg());
            if (!phi_pred_opt.has_value())
                continue;

            // τ (tau) = angular deviation from Bragg position to center of oscillation range
            float tau_rad = deg_to_rad(wrap_deg_pm180(phi_pred_opt.value()));

            // Calculate diffracted beam direction S = S0 + p (at diffracting condition)
            // For zeta calculation, we need S at the predicted phi angle
            const float phi_pred_rad = deg_to_rad(phi_pred_opt.value());
            const float cos_phi = std::cos(phi_pred_rad);
            const float sin_phi = std::sin(phi_pred_rad);

            // Rotate pstar by predicted phi around m2 axis
            const float p_m2 = pstar * m2;
            const Coord p_parallel = m2 * p_m2;
            const Coord p_perp = pstar - p_parallel;
            const Coord m2_cross_p = m2 % pstar;

            const Coord p_rotated = p_parallel + p_perp * cos_phi + m2_cross_p * sin_phi;
            const Coord S = S0 + p_rotated;

            // Calculate zeta (XDS convention)
            float zeta = calc_zeta(S, S0, m2);

            // Filter out reflections with very small zeta (poorly determined)
            if (zeta < 0.1f)
                continue;

            tau_values.push_back(tau_rad);
            zeta_values.push_back(zeta);
        }

        if (tau_values.size() < 10)
            return std::nullopt;

        // Find sigma_M by maximizing log-likelihood
        float sigma_M_rad = find_sigma_M_mle(tau_values, zeta_values, delta_phi_rad);

        return rad_to_deg(sigma_M_rad);
    }

    std::optional<float> CalcMosaicity(const DiffractionExperiment& experiment,
                                       const std::vector<SpotToSave> &spots,
                                       const Coord &astar, const Coord &bstar, const Coord &cstar) {
        const auto &axis = experiment.GetGoniometer();
        if (axis.has_value()) {
            const Coord w = axis->GetAxis().Normalize();
            const Coord S0 = experiment.GetScatteringVector();
            const float wedge_deg = axis->GetWedge_deg();

            double sum_sq = 0.0;
            int count = 0;

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

                const Coord pstar = 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(pstar, S0, w, 0.0, wedge_deg);
                if (!phi_pred_deg_opt.has_value())
                    continue;

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

            if (count > 0) {
                return static_cast<float>(std::sqrt(sum_sq / static_cast<double>(count)));
            }
        }
        return std::nullopt;
    }

} // namespace

bool AnalyzeIndexing(DataMessage &message,
                     const DiffractionExperiment &experiment,
                     const CrystalLattice &latt) {
    std::vector<uint8_t> indexed_spots(message.spots.size());

    // 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 bool index_ice_ring = experiment.GetIndexingSettings().GetIndexIceRings();
    const auto geom = experiment.GetDiffractionGeometry();
    const auto indexing_tolerance = experiment.GetIndexingSettings().GetTolerance();
    const auto indexing_tolerance_sq = indexing_tolerance * indexing_tolerance;
    const auto viable_cell_min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots();

    size_t nspots_ref = 0;
    size_t nspots_indexed = 0;

    // 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_sq) {
            if (index_ice_ring || !message.spots[i].ice_ring)
                nspots_indexed++;
            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);
        }
        if (index_ice_ring || !message.spots[i].ice_ring)
            nspots_ref++;
    }

    if (nspots_indexed >= viable_cell_min_spots && nspots_indexed >= std::lround(min_percentage_spots * nspots_ref)) {
        auto uc = latt.GetUnitCell();
        if (!ok(uc.a) || !ok(uc.b) || !ok(uc.c) || !ok(uc.alpha) || !ok(uc.beta) || !ok(uc.gamma))
            return false;

        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 = nspots_indexed;
        message.indexing_lattice = latt;
        message.indexing_unit_cell = latt.GetUnitCell();
        message.mosaicity_deg = CalcMosaicityXDS(experiment, message.spots, astar, bstar, cstar);

        return true;
    }

    message.indexing_result = false;
    return false;
}