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

#include "CellEsd.h"

#include <algorithm>
#include <cmath>

#include "../../common/JFJochMath.h"
#include "../../common/ParallelFor.h"

namespace {
    using Vec6 = Eigen::Matrix<double, 6, 1>;
    // At most six parameters: fixed-capacity, so the per-observation products never touch the heap.
    using VecK = Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 6, 1>;
    using MatK = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, 0, 6, 6>;

    // A symmetric matrix as (g11, g22, g33, g12, g13, g23) and back.
    Vec6 ToVec6(const Eigen::Matrix3d &m) {
        Vec6 v;
        v << m(0, 0), m(1, 1), m(2, 2), m(0, 1), m(0, 2), m(1, 2);
        return v;
    }

    Eigen::Matrix3d FromVec6(const Vec6 &v) {
        Eigen::Matrix3d m;
        m << v(0), v(3), v(4),
             v(3), v(1), v(5),
             v(4), v(5), v(2);
        return m;
    }

    // h^T G h = hh . (g11, g22, g33, g12, g13, g23)
    Vec6 IndexProducts(const CellObservation &o) {
        const double h = o.h, k = o.k, l = o.l;
        Vec6 v;
        v << h * h, k * k, l * l, 2 * h * k, 2 * h * l, 2 * k * l;
        return v;
    }

    double ToDeg(double rad) { return rad * 180.0 / PI; }
}

std::vector<Eigen::Matrix3d> MetricBasis(const gemmi::SpaceGroup &sg) {
    // Project each of the six symmetric unit matrices onto the metrics the group leaves invariant,
    // averaging R G R^T over the rotation parts, and keep an orthonormal basis of what comes out.
    const gemmi::GroupOps ops = sg.operations();
    std::vector<Eigen::Matrix3d> rots;
    for (const gemmi::Op &op : ops.sym_ops) {
        Eigen::Matrix3d r;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                r(i, j) = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
        rots.push_back(r);
    }
    Eigen::Matrix<double, 6, 6> projected;
    for (int m = 0; m < 6; m++) {
        Vec6 e = Vec6::Zero();
        e(m) = 1.0;
        const Eigen::Matrix3d g = FromVec6(e);
        Eigen::Matrix3d sum = Eigen::Matrix3d::Zero();
        for (const auto &r : rots)
            sum += r * g * r.transpose();
        projected.col(m) = ToVec6(sum / static_cast<double>(rots.size()));
    }
    Eigen::JacobiSVD<Eigen::Matrix<double, 6, 6>> svd(projected, Eigen::ComputeFullU);
    std::vector<Eigen::Matrix3d> basis;
    for (int m = 0; m < 6; m++)
        if (svd.singularValues()(m) > 1e-8 * svd.singularValues()(0))
            basis.push_back(FromVec6(svd.matrixU().col(m)));
    return basis;
}

std::array<double, 6> CellFromReciprocalMetric(const Eigen::Matrix3d &gstar) {
    const Eigen::Matrix3d g = gstar.inverse();
    const double a = std::sqrt(g(0, 0)), b = std::sqrt(g(1, 1)), c = std::sqrt(g(2, 2));
    return {a, b, c, ToDeg(std::acos(g(1, 2) / (b * c))), ToDeg(std::acos(g(0, 2) / (a * c))),
            ToDeg(std::acos(g(0, 1) / (a * b)))};
}

std::array<double, 6> CellEsdFromCovariance(const std::vector<Eigen::Matrix3d> &basis, const Eigen::VectorXd &p,
                                            const Eigen::MatrixXd &cov) {
    Eigen::Matrix3d gstar = Eigen::Matrix3d::Zero();
    for (size_t k = 0; k < basis.size(); k++)
        gstar += p(k) * basis[k];
    const Eigen::Matrix3d g = gstar.inverse();
    const double len[3] = {std::sqrt(g(0, 0)), std::sqrt(g(1, 1)), std::sqrt(g(2, 2))};
    // The angle opposite axis i lies between the other two, j and l: cos = G_jl / (len_j len_l).
    const int pair[3][2] = {{1, 2}, {0, 2}, {0, 1}};

    // Analytic derivatives of the six parameters with respect to each p_k: dG = -G dG* G.
    Eigen::MatrixXd jac(6, static_cast<Eigen::Index>(basis.size()));
    for (size_t k = 0; k < basis.size(); k++) {
        const Eigen::Matrix3d dg = -g * basis[k] * g;
        double dlen[3];
        for (int i = 0; i < 3; i++) {
            dlen[i] = dg(i, i) / (2.0 * len[i]);
            jac(i, k) = dlen[i];
        }
        for (int i = 0; i < 3; i++) {
            const int j = pair[i][0], l = pair[i][1];
            const double cosv = g(j, l) / (len[j] * len[l]);
            const double dcos = dg(j, l) / (len[j] * len[l]) - cosv * (dlen[j] / len[j] + dlen[l] / len[l]);
            jac(3 + i, k) = ToDeg(-dcos / std::sqrt(1.0 - cosv * cosv));
        }
    }
    const Eigen::MatrixXd cell_cov = jac * cov * jac.transpose();
    std::array<double, 6> esd{};
    // A parameter the symmetry fixes comes out at rounding level (1e-19), not at zero: make it zero.
    for (int i = 0; i < 6; i++) {
        esd[i] = std::sqrt(std::max(0.0, cell_cov(i, i)));
        if (esd[i] < 1e-12)
            esd[i] = 0.0;
    }
    return esd;
}

MetricFit FitReciprocalMetric(std::span<const CellObservation> obs, const std::vector<Eigen::Matrix3d> &basis) {
    const int nparam = static_cast<int>(basis.size());
    Eigen::Matrix<double, 6, Eigen::Dynamic, 0, 6, 6> v(6, nparam);
    for (int k = 0; k < nparam; k++)
        v.col(k) = ToVec6(basis[k]);

    MetricFit fit;
    fit.used.assign(obs.size(), true);
    MatK normal;
    double s2 = 0.0;
    // Fit, reject the observations further than 4 sigma off it, fit again - until nothing changes.
    for (int iter = 0; iter < 5; iter++) {
        normal = MatK::Zero(nparam, nparam);
        VecK rhs = VecK::Zero(nparam);
        int64_t n = 0;
        for (size_t i = 0; i < obs.size(); i++) {
            if (!fit.used[i])
                continue;
            const double q2 = static_cast<double>(obs[i].q) * obs[i].q;
            const VecK a = v.transpose() * IndexProducts(obs[i]);
            normal.noalias() += a * a.transpose() / q2;   // weight 1/q^2 on q^2: one e.s.d. on q
            rhs.noalias() += a;                            // a * q^2 / q^2
            n++;
        }
        if (n < 10 * nparam + 10)
            return {};
        fit.p = normal.ldlt().solve(rhs);
        double wr2 = 0.0;
        std::vector<double> z(obs.size());
        for (size_t i = 0; i < obs.size(); i++) {
            const double q2 = static_cast<double>(obs[i].q) * obs[i].q;
            const VecK a = v.transpose() * IndexProducts(obs[i]);
            const double r = q2 - a.dot(fit.p);
            z[i] = r * r / q2;
            if (fit.used[i])
                wr2 += z[i];
        }
        s2 = wr2 / static_cast<double>(n - nparam);
        bool changed = false;
        for (size_t i = 0; i < obs.size(); i++) {
            const bool keep = z[i] <= 16.0 * s2;
            changed |= keep != fit.used[i];
            fit.used[i] = keep;
        }
        if (!changed)
            break;
    }
    fit.cov = s2 * Eigen::MatrixXd(normal).inverse();
    fit.ok = fit.p.allFinite() && fit.cov.allFinite();
    return fit;
}

std::optional<CellEsd> EstimateCellEsd(std::span<const CellObservation> obs, const gemmi::SpaceGroup &sg,
                                       double wavelength_A, int nblocks) {
    const auto basis = MetricBasis(sg);
    const MetricFit fit = FitReciprocalMetric(obs, basis);
    if (!fit.ok)
        return {};
    const auto metric = [&basis](const Eigen::VectorXd &p) {
        Eigen::Matrix3d g = Eigen::Matrix3d::Zero();
        for (size_t k = 0; k < basis.size(); k++)
            g += p(k) * basis[k];
        return g;
    };

    CellEsd out;
    out.esd_lsq = CellEsdFromCovariance(basis, fit.p, fit.cov);
    const auto cell = CellFromReciprocalMetric(metric(fit.p));
    out.fitted = UnitCell{static_cast<float>(cell[0]), static_cast<float>(cell[1]), static_cast<float>(cell[2]),
                          static_cast<float>(cell[3]), static_cast<float>(cell[4]), static_cast<float>(cell[5])};

    std::vector<int64_t> keys;
    double q_min = INFINITY, q_max = 0.0;
    int32_t frame_min = INT32_MAX, frame_max = INT32_MIN;
    for (size_t i = 0; i < obs.size(); i++) {
        if (!fit.used[i])
            continue;
        const auto &o = obs[i];
        keys.push_back(((static_cast<int64_t>(o.h) + 1024) << 22) | ((static_cast<int64_t>(o.k) + 1024) << 11)
                       | (static_cast<int64_t>(o.l) + 1024));
        q_min = std::min(q_min, static_cast<double>(o.q));
        q_max = std::max(q_max, static_cast<double>(o.q));
        frame_min = std::min(frame_min, o.frame);
        frame_max = std::max(frame_max, o.frame);
    }
    out.observations = static_cast<int64_t>(keys.size());
    std::sort(keys.begin(), keys.end());
    out.reflections = std::unique(keys.begin(), keys.end()) - keys.begin();
    out.theta_min_deg = ToDeg(std::asin(std::min(1.0, wavelength_A * q_min / 2.0)));
    out.theta_max_deg = ToDeg(std::asin(std::min(1.0, wavelength_A * q_max / 2.0)));

    // The same fit on blocks of consecutive frames: the scatter of their cells, over sqrt(blocks).
    // Least-squares cell e.s.d.s alone are too small: following Herbstein (2000) Acta Cryst. B56, 547-557
    std::vector<std::array<double, 6>> block_cells;
    const int64_t span = static_cast<int64_t>(frame_max) - frame_min + 1;
    std::vector<std::vector<CellObservation>> block_obs(nblocks);
    for (size_t i = 0; i < obs.size(); i++)
        if (fit.used[i])
            block_obs[(obs[i].frame - frame_min) * nblocks / span].push_back(obs[i]);
    for (const auto &b : block_obs) {
        const MetricFit bf = FitReciprocalMetric(b, basis);
        if (bf.ok)
            block_cells.push_back(CellFromReciprocalMetric(metric(bf.p)));
    }
    out.blocks = static_cast<int>(block_cells.size());
    out.esd = out.esd_lsq;
    if (out.blocks >= 4) {
        std::array<double, 6> esd{};
        for (int j = 0; j < 6; j++) {
            double mean = 0.0;
            for (const auto &c : block_cells)
                mean += c[j] / out.blocks;
            double var = 0.0;
            for (const auto &c : block_cells)
                var += (c[j] - mean) * (c[j] - mean) / (out.blocks - 1);
            esd[j] = std::sqrt(var / out.blocks);
            out.esd[j] = std::max(out.esd[j], esd[j]);
        }
        out.esd_blocks = esd;
    }
    return out;
}

std::vector<CellObservation> GatherCellObservations(const std::vector<IntegrationOutcome> &outcomes,
                                                    size_t nthreads) {
    std::vector<std::vector<CellObservation>> per_frame(outcomes.size());
    ParallelFor(static_cast<int>(outcomes.size()), nthreads, [&](int f) {
        const IntegrationOutcome &out = outcomes[f];
        for (const auto &r : out.reflections) {
            if (!r.observed || r.clipped || r.overloaded || r.on_ice_ring || !(r.sigma > 0.0f)
                || !(r.I > 3.0f * r.sigma) || !std::isfinite(r.observed_x) || !std::isfinite(r.observed_y))
                continue;
            const float q = out.geom.DetectorToRecip(r.observed_x, r.observed_y).Length();
            if (q > 0.0f)
                per_frame[f].push_back(CellObservation{r.h, r.k, r.l, q, f});
        }
    });
    std::vector<CellObservation> obs;
    for (const auto &v : per_frame)
        obs.insert(obs.end(), v.begin(), v.end());
    return obs;
}