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

#include "RefinementADP.h"

#include <algorithm>
#include <cmath>
#include <limits>

#include <Eigen/Dense>

namespace {

constexpr double DEG = 3.14159265358979323846 / 180.0;

}  // namespace

gemmi::SMat33<double> TlsBody::U(const gemmi::Position &pos) const {
    const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
    const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
    const gemmi::Mat33 ala = a.multiply(L.as_mat33()).multiply(a.transpose());
    const gemmi::Mat33 as = a.multiply(S);
    gemmi::SMat33<double> u = T;
    u.u11 += ala[0][0] + 2 * as[0][0];
    u.u22 += ala[1][1] + 2 * as[1][1];
    u.u33 += ala[2][2] + 2 * as[2][2];
    u.u12 += ala[0][1] + as[0][1] + as[1][0];
    u.u13 += ala[0][2] + as[0][2] + as[2][0];
    u.u23 += ala[1][2] + as[1][2] + as[2][1];
    return u;
}

std::array<double, TlsBody::PARAMETERS> TlsBody::Get() const {
    return {T.u11, T.u22, T.u33, T.u12, T.u13, T.u23, L.u11, L.u22, L.u33, L.u12, L.u13, L.u23,
            S[0][0], S[1][1], S[0][1], S[0][2], S[1][0], S[1][2], S[2][0], S[2][1]};
}

void TlsBody::Set(const std::array<double, PARAMETERS> &p) {
    T = {p[0], p[1], p[2], p[3], p[4], p[5]};
    L = {p[6], p[7], p[8], p[9], p[10], p[11]};
    S = gemmi::Mat33(p[12], p[14], p[15], p[16], p[13], p[17], p[18], p[19], -(p[12] + p[13]));
}

std::array<std::array<double, 6>, TlsBody::PARAMETERS> TlsBody::Derivatives(const gemmi::Position &pos) const {
    std::array<std::array<double, 6>, PARAMETERS> d{};
    TlsBody unit;
    unit.origin = origin;
    for (int k = 0; k < PARAMETERS; k++) {
        std::array<double, PARAMETERS> p{};
        p[k] = 1;
        unit.Set(p);
        const gemmi::SMat33<double> u = unit.U(pos);
        d[k] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
    }
    return d;
}

std::array<double, TlsBody::PARAMETERS> TlsBody::Gradient(const std::vector<gemmi::Atom *> &atom_list,
                                                         const std::vector<std::array<double, 6>> &grad_u) const {
    gemmi::Mat33 mt(0), ml(0), ms(0);
    for (int i : atoms) {
        const std::array<double, 6> &g = grad_u[i];
        const gemmi::Mat33 gm(g[0], g[3] / 2, g[4] / 2, g[3] / 2, g[1], g[5] / 2, g[4] / 2, g[5] / 2, g[2]);
        const gemmi::Position &pos = atom_list[i]->pos;
        const double x = pos.x - origin.x, y = pos.y - origin.y, z = pos.z - origin.z;
        const gemmi::Mat33 a(0, z, -y, -z, 0, x, y, -x, 0);
        const gemmi::Mat33 atg = a.transpose().multiply(gm);
        const gemmi::Mat33 atga = atg.multiply(a);
        for (int u = 0; u < 3; u++)
            for (int v = 0; v < 3; v++) {
                mt[u][v] += gm[u][v];
                ml[u][v] += atga[u][v];
                ms[u][v] += 2 * atg[u][v];
            }
    }
    // symmetric T and L: an off-diagonal parameter takes both of its elements; S33 = -(S11 + S22)
    return {mt[0][0], mt[1][1], mt[2][2], 2 * mt[0][1], 2 * mt[0][2], 2 * mt[1][2],
            ml[0][0], ml[1][1], ml[2][2], 2 * ml[0][1], 2 * ml[0][2], 2 * ml[1][2],
            ms[0][0] - ms[2][2], ms[1][1] - ms[2][2], ms[0][1], ms[0][2], ms[1][0], ms[1][2], ms[2][0], ms[2][1]};
}

std::vector<TlsBody> TlsBodiesFromModel(const gemmi::Model &model, const std::vector<gemmi::TlsGroup> &groups) {
    std::vector<TlsBody> bodies;
    for (const gemmi::TlsGroup &g : groups) {
        TlsBody b;
        b.id = g.id;
        b.origin = g.origin;
        b.T = g.T;
        b.L = g.L.scaled(DEG * DEG);
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                b.S[i][j] = g.S[i][j] * DEG;
        // the trace of S does not enter U; S33 is taken as -(S11 + S22)
        const double t = (b.S[0][0] + b.S[1][1] + b.S[2][2]) / 3.0;
        for (int i = 0; i < 3; i++)
            b.S[i][i] -= t;
        bool finite = std::isfinite(b.origin.x) && std::isfinite(b.origin.y) && std::isfinite(b.origin.z);
        for (double v : b.Get())
            finite = finite && std::isfinite(v);
        if (finite)
            bodies.push_back(b);
    }
    // the atoms of each group's residues, an atom to the first group that names it
    int i = 0;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues) {
            int body = -1;
            if (!r.is_water())
                for (size_t k = 0; k < groups.size() && body < 0; k++) {
                    for (const gemmi::TlsGroup::Selection &sel : groups[k].selections)
                        if (sel.chain == ch.name && !(r.seqid < sel.res_begin) && !(sel.res_end < r.seqid))
                            body = static_cast<int>(k);
                }
            // the group index among those kept
            int kept = -1;
            if (body >= 0)
                for (size_t k = 0; k < bodies.size(); k++)
                    if (bodies[k].id == groups[body].id)
                        kept = static_cast<int>(k);
            for (size_t a = 0; a < r.atoms.size(); a++, i++)
                if (kept >= 0)
                    bodies[kept].atoms.push_back(i);
        }
    std::erase_if(bodies, [](const TlsBody &b) { return b.atoms.empty(); });
    return bodies;
}

std::vector<gemmi::TlsGroup> TlsBodiesToGroups(const gemmi::Model &model, const std::vector<TlsBody> &bodies) {
    // per atom its body, per residue the body of its first atom in one
    std::vector<int> atom_body;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            atom_body.insert(atom_body.end(), r.atoms.size(), -1);
    for (size_t b = 0; b < bodies.size(); b++)
        for (int a : bodies[b].atoms)
            atom_body[a] = static_cast<int>(b);
    std::vector<gemmi::TlsGroup> groups(bodies.size());
    for (size_t b = 0; b < bodies.size(); b++) {
        gemmi::TlsGroup &g = groups[b];
        g.id = std::to_string(b + 1);
        g.num_id = static_cast<short>(b + 1);
        g.origin = bodies[b].origin;
        g.T = bodies[b].T;
        g.L = bodies[b].L.scaled(1.0 / (DEG * DEG));
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                g.S[i][j] = bodies[b].S[i][j] / DEG;
    }
    size_t i = 0;
    for (const gemmi::Chain &ch : model.chains) {
        int open = -1;   // the body whose run of residues is open in this chain
        for (const gemmi::Residue &r : ch.residues) {
            int body = -1;
            for (size_t a = 0; a < r.atoms.size(); a++)
                if (atom_body[i + a] >= 0 && body < 0)
                    body = atom_body[i + a];
            i += r.atoms.size();
            if (body >= 0 && body == open)
                groups[body].selections.back().res_end = r.seqid;
            else if (body >= 0) {
                groups[body].selections.push_back({ch.name, r.seqid, r.seqid, ""});
            }
            open = body;
        }
    }
    return groups;
}

URestraints MakeURestraints(const std::vector<std::pair<int, int>> &b_similarity_pairs, const std::vector<char> &refines_u) {
    URestraints r;
    std::vector<char> paired(refines_u.size(), 0);
    for (const auto &[i, j] : b_similarity_pairs)
        if (refines_u[i] && refines_u[j]) {
            r.pairs.emplace_back(i, j);
            paired[i] = paired[j] = 1;
        }
    for (size_t i = 0; i < refines_u.size(); i++)
        if (refines_u[i] && !paired[i])
            r.isotropic.push_back(static_cast<int>(i));
    return r;
}

double EvaluateURestraints(const URestraints &r, const std::vector<gemmi::Position> &pos,
                           const std::vector<std::array<double, 6>> &u, std::vector<std::array<double, 6>> *grad) {
    // a symmetric matrix as its six elements: the off-diagonal ones count twice in a Frobenius norm, and a derivative
    // with respect to U_ab (a != b) is the sum of the matrix elements ab and ba
    double e = 0;
    const double wr = 1.0 / (URestraints::SIGMA_RIGID * URestraints::SIGMA_RIGID);
    const double ws = 1.0 / (URestraints::SIGMA_SIMILAR * URestraints::SIGMA_SIMILAR);
    const double wi = 1.0 / (URestraints::SIGMA_ISOTROPIC * URestraints::SIGMA_ISOTROPIC);
    for (const auto &[i, j] : r.pairs) {
        std::array<double, 6> d;
        for (int c = 0; c < 6; c++)
            d[c] = u[i][c] - u[j][c];
        gemmi::Vec3 n = pos[j] - pos[i];
        const double len = n.length();
        std::array<double, 6> g{};
        if (len > 0) {
            n = n / len;
            // v = (U_i - U_j) n
            const double v[3] = {d[0] * n.x + d[3] * n.y + d[4] * n.z, d[3] * n.x + d[1] * n.y + d[5] * n.z,
                                 d[4] * n.x + d[5] * n.y + d[2] * n.z};
            const double nn[3] = {n.x, n.y, n.z};
            e += wr * (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]);
            g[0] += 2 * wr * v[0] * nn[0];
            g[1] += 2 * wr * v[1] * nn[1];
            g[2] += 2 * wr * v[2] * nn[2];
            g[3] += 2 * wr * (v[0] * nn[1] + v[1] * nn[0]);
            g[4] += 2 * wr * (v[0] * nn[2] + v[2] * nn[0]);
            g[5] += 2 * wr * (v[1] * nn[2] + v[2] * nn[1]);
        }
        e += ws * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
        for (int c = 0; c < 3; c++)
            g[c] += 2 * ws * d[c];
        for (int c = 3; c < 6; c++)
            g[c] += 4 * ws * d[c];
        if (grad != nullptr)
            for (int c = 0; c < 6; c++) {
                (*grad)[i][c] += g[c];
                (*grad)[j][c] -= g[c];
            }
    }
    for (int i : r.isotropic) {
        const double t = (u[i][0] + u[i][1] + u[i][2]) / 3.0;
        const double d[6] = {u[i][0] - t, u[i][1] - t, u[i][2] - t, u[i][3], u[i][4], u[i][5]};
        e += wi * (d[0] * d[0] + d[1] * d[1] + d[2] * d[2] + 2 * (d[3] * d[3] + d[4] * d[4] + d[5] * d[5]));
        if (grad != nullptr) {
            // d/dU_aa of the traceless part's norm: 2 d_aa (the trace's own derivative sums to zero against d)
            for (int c = 0; c < 3; c++)
                (*grad)[i][c] += 2 * wi * d[c];
            for (int c = 3; c < 6; c++)
                (*grad)[i][c] += 4 * wi * d[c];
        }
    }
    return e;
}

std::vector<int> TlsSegments(const std::vector<std::vector<int>> &residue_atoms, const std::vector<gemmi::Position> &pos,
                             const std::vector<double> &b, int k, int min_residues) {
    using Mat = Eigen::Matrix<double, 10, 10>;
    using Vec = Eigen::Matrix<double, 10, 1>;
    const int n = static_cast<int>(residue_atoms.size());
    k = std::min(k, n / std::max(min_residues, 1));
    if (k <= 1)
        return {0};
    // The chain in blocks of consecutive residues (cuts between blocks only), so that a long chain stays a few hundred
    // blocks: the dynamic programming is quadratic in them.
    const int per_block = std::max(1, (n + 299) / 300);
    const int min_blocks = std::max(1, (min_residues + per_block - 1) / per_block);
    std::vector<int> block_first;
    for (int r = 0; r < n; r += per_block)
        block_first.push_back(r);
    const int nb = static_cast<int>(block_first.size());
    k = std::min(k, nb / min_blocks);
    if (k <= 1)
        return {0};
    // the normal equations of each block, about the chain's centroid in units of 10 A (for their conditioning)
    gemmi::Position c;
    int count = 0;
    for (const std::vector<int> &ra : residue_atoms)
        for (int a : ra) {
            c += pos[a];
            count++;
        }
    if (count == 0)
        return {0};
    c /= count;
    std::vector<Mat> ata(nb + 1, Mat::Zero());
    std::vector<Vec> atb(nb + 1, Vec::Zero());
    std::vector<double> bb(nb + 1, 0.0);
    for (int q = 0; q < nb; q++) {
        ata[q + 1] = ata[q];
        atb[q + 1] = atb[q];
        bb[q + 1] = bb[q];
        const int last = q + 1 < nb ? block_first[q + 1] : n;
        for (int r = block_first[q]; r < last; r++)
            for (int a : residue_atoms[r]) {
                const double x = (pos[a].x - c.x) / 10, y = (pos[a].y - c.y) / 10, z = (pos[a].z - c.z) / 10;
                Vec f;
                f << 1, x, y, z, x * x, y * y, z * z, x * y, x * z, y * z;
                ata[q + 1] += f * f.transpose();
                atb[q + 1] += b[a] * f;
                bb[q + 1] += b[a] * b[a];
            }
    }
    // the residual of the fit of blocks [i, j)
    auto cost = [&](int i, int j) {
        const Mat m = ata[j] - ata[i] + 1e-6 * Mat::Identity();
        const Vec v = atb[j] - atb[i];
        const Vec x = m.ldlt().solve(v);
        return std::max(0.0, bb[j] - bb[i] - v.dot(x));
    };
    const double INF = std::numeric_limits<double>::infinity();
    // best[s][j]: blocks [0, j) in s + 1 segments; from[s][j]: where the last of them starts
    std::vector<std::vector<double>> best(k, std::vector<double>(nb + 1, INF));
    std::vector<std::vector<int>> from(k, std::vector<int>(nb + 1, 0));
    std::vector<std::vector<double>> seg(nb + 1, std::vector<double>(nb + 1, INF));
    for (int i = 0; i < nb; i++)
        for (int j = i + min_blocks; j <= nb; j++)
            seg[i][j] = cost(i, j);
    for (int j = min_blocks; j <= nb; j++)
        best[0][j] = seg[0][j];
    for (int s = 1; s < k; s++)
        for (int j = (s + 1) * min_blocks; j <= nb; j++)
            for (int i = s * min_blocks; i <= j - min_blocks; i++)
                if (best[s - 1][i] + seg[i][j] < best[s][j]) {
                    best[s][j] = best[s - 1][i] + seg[i][j];
                    from[s][j] = i;
                }
    std::vector<int> starts;
    int j = nb;
    for (int s = k - 1; s > 0; s--) {
        j = from[s][j];
        starts.push_back(block_first[j]);
    }
    starts.push_back(0);
    std::reverse(starts.begin(), starts.end());
    return starts;
}