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

#include "CrystalSetting.h"

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <map>

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

namespace {
    constexpr double DEG = PI / 180.0;

    // The origin shifts every crystallographic setting change lies on, the common ones first.
    constexpr int TWELFTHS[] = {0, 6, 3, 9, 4, 8, 2, 10, 1, 5, 7, 11};

    // Is `g` the group whose sorted operator list is `target`, on some origin? The rotations and the
    // centring are compared first - they do not depend on the origin - so the shift search only runs
    // for a candidate that can match.
    bool SameGroupUpToOrigin(const gemmi::GroupOps &g, const gemmi::GroupOps &target,
                             const std::vector<gemmi::Op> &target_sorted) {
        if (g.sym_ops.size() != target.sym_ops.size() || !g.has_same_centring(target))
            return false;
        for (const gemmi::Op &op : g.sym_ops)
            if (target.find_by_rotation(op.rot) == nullptr)
                return false;
        for (int i : TWELFTHS)
            for (int j : TWELFTHS)
                for (int k : TWELFTHS) {
                    gemmi::Op shift = gemmi::Op::identity();
                    shift.tran = {i * gemmi::Op::DEN / 12, j * gemmi::Op::DEN / 12, k * gemmi::Op::DEN / 12};
                    gemmi::GroupOps s = g;
                    s.change_basis_forward(shift);
                    if (s.all_ops_sorted() == target_sorted)
                        return true;
                }
        return false;
    }

    // How far a change of basis is from leaving everything where it is: the entries that differ from
    // the identity, then the negative ones. Only ever used to break a tie between equal candidates.
    int IdentityDistance(const gemmi::Op &cob) {
        int changed = 0, negative = 0;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++) {
                const int e = cob.rot[i][j] / gemmi::Op::DEN;
                changed += e != (i == j ? 1 : 0);
                negative += e < 0;
            }
        return 10 * changed + negative;
    }
}

gemmi::Op HklOperator(const gemmi::Op &cob) {
    gemmi::Op op = cob.inverse();
    op.tran = {0, 0, 0};
    return op;
}

std::string IndexTriplet(const gemmi::Op &hkl_op) {
    std::string out;
    for (int i = 0; i < 3; i++) {
        std::string term;
        for (int j = 0; j < 3; j++) {
            const int e = hkl_op.rot[j][i] / gemmi::Op::DEN;   // h'_i = sum_j rot[j][i] h_j
            if (e == 0)
                continue;
            if (e < 0)
                term += "-";
            else if (!term.empty())
                term += "+";
            if (std::abs(e) != 1)
                term += std::to_string(std::abs(e));
            term += "hkl"[j];
        }
        out += (i ? "," : "") + (term.empty() ? std::string("0") : term);
    }
    return out;
}

gemmi::Mat33 BasisMatrix(const gemmi::Op &cob) {
    // h' = P h and apply_to_hkl computes rot^T h, so P is the transpose of the index operator.
    const gemmi::Op hkl = HklOperator(cob);
    gemmi::Mat33 p;
    for (int i = 0; i < 3; i++)
        for (int j = 0; j < 3; j++)
            p.a[i][j] = static_cast<double>(hkl.rot[j][i]) / gemmi::Op::DEN;
    return p;
}

UnitCell CellInBasis(const UnitCell &cell, const gemmi::Op &cob) {
    gemmi::UnitCell g = cell;
    gemmi::Op rot_only = cob;
    rot_only.tran = {0, 0, 0};
    const gemmi::UnitCell c = g.changed_basis_forward(rot_only, false);
    return UnitCell{static_cast<float>(c.a), static_cast<float>(c.b), static_cast<float>(c.c),
                    static_cast<float>(c.alpha), static_cast<float>(c.beta), static_cast<float>(c.gamma)};
}

const gemmi::SpaceGroup *SpaceGroupInBasis(const gemmi::SpaceGroup &sg, const gemmi::Op &cob) {
    for (int i : TWELFTHS)
        for (int j : TWELFTHS)
            for (int k : TWELFTHS) {
                gemmi::Op shifted = cob;
                shifted.tran = {i * gemmi::Op::DEN / 12, j * gemmi::Op::DEN / 12, k * gemmi::Op::DEN / 12};
                gemmi::GroupOps gops = sg.operations();
                gops.change_basis_forward(shifted);
                if (const gemmi::SpaceGroup *found = gemmi::find_spacegroup_by_ops(gops))
                    return found;
            }
    return nullptr;
}

double MetricViolation(const UnitCell &uc, const gemmi::SpaceGroup &sg) {
    const double a = uc.a, b = uc.b, c = uc.c;
    const double ab = a * b * std::cos(uc.gamma * DEG);
    const double ac = a * c * std::cos(uc.beta * DEG);
    const double bc = b * c * std::cos(uc.alpha * DEG);
    const double g[3][3] = {{a * a, ab, ac}, {ab, b * b, bc}, {ac, bc, c * c}};
    const double scale = std::max({a * a, b * b, c * c});
    if (!(scale > 0.0))
        return 0.0;
    double worst = 0.0;
    for (const gemmi::Op &op : sg.operations()) {
        double r[3][3];
        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;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++) {
                double t = 0.0;
                for (int k = 0; k < 3; k++)
                    for (int l = 0; l < 3; l++)
                        t += r[k][i] * g[k][l] * r[l][j];
                worst = std::max(worst, std::fabs(t - g[i][j]) / scale);
            }
    }
    return worst;
}

bool CellsCorrespond(const gemmi::UnitCell &a, const gemmi::UnitCell &b) {
    auto len = [](double x, double y) { return std::fabs(x - y) <= 0.05 * std::max(x, y); };
    auto ang = [](double x, double y) { return std::fabs(x - y) <= 3.0; };
    return len(a.a, b.a) && len(a.b, b.b) && len(a.c, b.c)
           && ang(a.alpha, b.alpha) && ang(a.beta, b.beta) && ang(a.gamma, b.gamma);
}

const std::vector<gemmi::Op> &UnimodularOperators() {
    static const std::vector<gemmi::Op> ops = [] {
        std::vector<gemmi::Op> out{gemmi::Op::identity()};
        for (int pattern = 0; pattern < 19683; pattern++) {   // 3^9 matrices over {-1,0,1}
            int e[9], v = pattern;
            for (int &x : e) { x = v % 3 - 1; v /= 3; }
            const int det = e[0] * (e[4] * e[8] - e[5] * e[7])
                          - e[1] * (e[3] * e[8] - e[5] * e[6])
                          + e[2] * (e[3] * e[7] - e[4] * e[6]);
            if (det != 1)
                continue;
            gemmi::Op op = gemmi::Op::identity();
            for (int i = 0; i < 3; i++)
                for (int j = 0; j < 3; j++)
                    op.rot[i][j] = e[3 * i + j] * gemmi::Op::DEN;
            if (!(op == gemmi::Op::identity()))
                out.push_back(op);
        }
        return out;
    }();
    return ops;
}

std::vector<gemmi::Op> CellMappingOperators(const gemmi::UnitCell &from, const gemmi::UnitCell &to) {
    std::vector<gemmi::Op> out;
    for (const gemmi::Op &op : UnimodularOperators()) {
        gemmi::UnitCell moved = from;   // changed_basis_forward is not const
        if (CellsCorrespond(moved.changed_basis_forward(op, false), to))
            out.push_back(op);
    }
    return out;
}

gemmi::Op ChooseOutputSetting(const UnitCell &cell, const gemmi::SpaceGroup &sg,
                              const gemmi::SpaceGroup *target, const std::optional<UnitCell> &target_cell) {
    // A triclinic cell has no symmetry axis to name; the reduced cell the run indexed is its standard.
    if (!target_cell && sg.number <= 2)
        return gemmi::Op::identity();

    // The settings a candidate may land on: the one asked for, the reference one, or - to follow a
    // cell given on its own - any setting of the group.
    std::vector<const gemmi::SpaceGroup *> settings;
    if (target != nullptr)
        settings.push_back(target);
    else if (!target_cell)
        settings.push_back(gemmi::find_spacegroup_by_number(sg.number));
    else
        for (const gemmi::SpaceGroup &s : gemmi::spacegroup_tables::main)
            if (s.number == sg.number)
                settings.push_back(&s);
    std::vector<gemmi::GroupOps> setting_ops;
    std::vector<std::vector<gemmi::Op>> setting_sorted;
    for (const gemmi::SpaceGroup *s : settings) {
        setting_ops.push_back(s->operations());
        setting_sorted.push_back(setting_ops.back().all_ops_sorted());
    }

    const bool monoclinic = sg.crystal_system() == gemmi::CrystalSystem::Monoclinic;
    const bool orthorhombic = sg.crystal_system() == gemmi::CrystalSystem::Orthorhombic;
    const auto obliquity = [](const UnitCell &c) {
        return std::max({std::fabs(c.alpha - 90.0), std::fabs(c.beta - 90.0), std::fabs(c.gamma - 90.0)});
    };
    const auto distance = [&](const UnitCell &c) {
        const UnitCell &t = *target_cell;
        return std::fabs(c.a - t.a) / t.a + std::fabs(c.b - t.b) / t.b + std::fabs(c.c - t.c) / t.c
               + (std::fabs(c.alpha - t.alpha) + std::fabs(c.beta - t.beta) + std::fabs(c.gamma - t.gamma)) / 180.0;
    };
    // Is candidate x a better choice than y?
    const auto better = [&](const gemmi::Op &x, const UnitCell &cx, const gemmi::Op &y, const UnitCell &cy) {
        if (target_cell) {
            const double dx = distance(cx), dy = distance(cy);
            if (std::fabs(dx - dy) > 1e-6)
                return dx < dy;
        } else if (monoclinic) {
            const double ox = obliquity(cx), oy = obliquity(cy);
            if (std::fabs(ox - oy) > 0.05)
                return ox < oy;
        } else if (orthorhombic) {
            const double lx[3] = {cx.a, cx.b, cx.c}, ly[3] = {cy.a, cy.b, cy.c};
            for (int i = 0; i < 3; i++)
                if (std::fabs(lx[i] - ly[i]) > 1e-3 * std::max(lx[i], ly[i]))
                    return lx[i] < ly[i];
        }
        return IdentityDistance(x) < IdentityDistance(y);
    };

    const gemmi::GroupOps ops = sg.operations();
    std::optional<gemmi::Op> best;
    UnitCell best_cell{};
    for (const gemmi::Op &cob : UnimodularOperators()) {
        const UnitCell c = CellInBasis(cell, cob);
        if (target_cell && !CellsCorrespond(c, *target_cell))
            continue;
        // beta >= 90 (or whichever angle the unique axis leaves free), as ITA writes it.
        if (!target_cell && monoclinic && std::min({c.alpha, c.beta, c.gamma}) < 90.0 - 1e-3)
            continue;
        gemmi::GroupOps g = ops;
        g.change_basis_forward(cob);
        bool lands = false;
        for (size_t i = 0; i < settings.size() && !lands; i++)
            lands = SameGroupUpToOrigin(g, setting_ops[i], setting_sorted[i]);
        if (!lands)
            continue;
        if (!best || better(cob, c, *best, best_cell)) {
            best = cob;
            best_cell = c;
        }
    }
    // A cell given with -C that no description of this lattice comes near says nothing about axes.
    if (!best && target_cell)
        return ChooseOutputSetting(cell, sg, target, std::nullopt);
    return best.value_or(gemmi::Op::identity());
}

gemmi::Op SeatGroupByAbsences(const UnitCell &cell, const gemmi::SpaceGroup &sg,
                              const std::vector<ReflectionZ> &reflections) {
    const gemmi::GroupOps g = sg.operations();
    // The centring is the lattice's and is the same for every candidate, so its absences say nothing
    // about which axis is which.
    gemmi::GroupOps lattice;
    lattice.sym_ops = {gemmi::Op::identity()};
    lattice.cen_ops = g.cen_ops;

    // The candidates: the group as the current axes would see it, were the reflections put through each
    // change of basis that keeps a cell the group can describe and the centring the lattice has.
    std::vector<gemmi::Op> cobs;
    std::vector<gemmi::GroupOps> seen_here;
    for (const gemmi::Op &cob : UnimodularOperators()) {
        if (MetricViolation(CellInBasis(cell, cob), sg) > MAX_METRIC_VIOLATION)
            continue;
        gemmi::GroupOps here = g;
        here.change_basis_backward(cob);
        if (!here.has_same_centring(g))
            continue;
        cobs.push_back(cob);
        seen_here.push_back(here);
    }
    if (cobs.size() < 2)
        return gemmi::Op::identity();

    // Every reflection some candidate predicts absent falls in a class of reflections that the same
    // candidates predict absent (a parity class of one axial row, say). Each class is read as the ratio
    // of its mean I/sigma to that of the reflections no candidate predicts absent: near 0 absent, near
    // 1 present. A candidate is charged the ratio for each class it calls absent and one minus it for
    // each class it calls present. A class nothing measured is not a class here, so it charges no
    // candidate - a screw on a row the sweep never recorded is neither for nor against.
    std::map<std::vector<bool>, std::pair<double, int>> classes;
    double control_sum = 0.0;
    int control_n = 0;
    for (const auto &r : reflections) {
        if (lattice.is_systematically_absent(r.hkl))
            continue;
        std::vector<bool> absent_in(cobs.size());
        bool any = false;
        for (size_t i = 0; i < cobs.size(); i++) {
            absent_in[i] = seen_here[i].is_systematically_absent(r.hkl);
            any = any || absent_in[i];
        }
        if (any) {
            auto &c = classes[absent_in];
            c.first += r.z;
            c.second++;
        } else {
            control_sum += r.z;
            control_n++;
        }
    }
    if (classes.empty() || control_n == 0 || control_sum <= 0.0)
        return gemmi::Op::identity();
    const double control = control_sum / control_n;

    size_t best = 0;
    double best_cost = 0.0;
    for (size_t i = 0; i < cobs.size(); i++) {
        double cost = 0.0;
        for (const auto &[absent_in, sum_n] : classes) {
            const double ratio = std::clamp(sum_n.first / sum_n.second / control, 0.0, 1.0);
            cost += absent_in[i] ? ratio : 1.0 - ratio;
        }
        if (i == 0 || cost < best_cost) {
            best = i;
            best_cost = cost;
        }
    }
    return cobs[best];
}