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

#include "../../common/JFJochMath.h"
#include "LatticeSearch.h"
#include <gemmi/cellred.hpp>

#include <cmath>
#include <optional>

// How close the reduced beta has to be to 90 degrees for the two Niggli types to be genuinely
// interchangeable (see the retry at the end of LatticeSearch). Not the angle tolerance: that is how
// far a metric may sit from an ideal one and still be called it, which is far too generous here - a
// cell 2 degrees off the boundary is a real type-1 cell, and presenting it in the obtuse setting
// promotes a general triclinic lattice to C-centred monoclinic on residuals of ~2 degrees. Measured:
// the crystal this was found on sits 0.07 degrees from the boundary and matches on 0.006 to 0.135;
// the triclinic cell that must not be promoted sits 2.0 degrees from it.
constexpr double NIGGLI_TYPE_BOUNDARY_DEG = 0.5;

struct NiggliClass {
    int number;
    int type;
    bool cond_AB;
    bool cond_BC;
    double cond_D;
    double cond_E;
    double cond_F;
    bool cond_DEF;
    bool cond_2DF;
    gemmi::Mat33 reindex;
    gemmi::CrystalSystem system;
    char centering;
};

namespace {
// The body of both entry points. only_class, when given, keeps just the characters of that Bravais
// class - see LatticeSearchForClass. With no filter this is the original walk unchanged, and the
// triclinic character fits every metric, so it always returns a result.
std::optional<LatticeSearchResult> SearchCharacters(const CrystalLattice &L, double dist_tolerance,
                                                    double angle_tolerance,
                                                    const std::pair<gemmi::CrystalSystem, char> *only_class) {
    UnitCell uc = L.GetUnitCell();
    gemmi::UnitCell g_uc(uc.a, uc.b, uc.c, uc.alpha, uc.beta, uc.gamma);
    gemmi::GruberVector g_vec(g_uc, 'P', true);
    g_vec.niggli_reduce();

    CrystalLattice L_niggli = L;
    if (g_vec.change_of_basis)
        L_niggli = L.Multiply(gemmi::rot_as_mat33(g_vec.change_of_basis->rot).transpose());

    double A = g_vec.A;
    double B = g_vec.B;
    double C = g_vec.C;
    double D = g_vec.xi / 2;
    double E = g_vec.eta / 2;
    double F = g_vec.zeta / 2;

    // D, E, F are parameters so the table can also be built for the type-flipped setting below.
    auto make_classes = [&](double D, double E, double F) {
      return std::vector<NiggliClass>{
        {
            1, 1,
            true, true, A / 2, A / 2, A / 2, false, false,
            gemmi::Mat33{1, -1, 1, 1, 1, -1, -1, 1, 1},
            gemmi::CrystalSystem::Cubic, 'F'
        },
        {
            2, 1,
            true, true, D, D, D, false, false,
            {1, -1, 0, -1, 0, 1, -1, -1, -1},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            3, 2,
            true, true, 0, 0, 0, false, false,
            gemmi::Mat33{1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Cubic, 'P'
        },
        {
            5, 2,
            true, true, -A / 3, -A / 3, -A / 3, false, false,
            gemmi::Mat33{1, 0, 1, 1, 1, 0, 0, 1, 1},
            gemmi::CrystalSystem::Cubic, 'I'
        },
        {
            4, 2,
            true, true, D, D, D, false, false,
            {1, -1, 0, -1, 0, 1, -1, -1, -1},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            6, 2,
            true, true, D, D, F, true, false,
            {0, 1, 1, 1, 0, 1, 1, 1, 0},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            7, 2,
            true, true, D, E, E, true, false,
            {1, 0, 1, 1, 1, 0, 0, 1, 1},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            8, 2,
            true, true, D, E, F, true, false,
            {-1, -1, 0, -1, 0, -1, 0, -1, -1},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            9, 1,
            true, false, A / 2, A / 2, A / 2, false, false,
            {1, 0, 0, -1, 1, 0, -1, -1, 3},
            gemmi::CrystalSystem::Trigonal, 'R'
        },

        {
            10, 1,
            true, false, D, D, F, false, false,
            {1, 1, 0, 1, -1, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            11, 2,
            true, false, 0, 0, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Tetragonal, 'P'
        },
        {
            12, 2,
            true, false, 0, 0, -A / 2, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Hexagonal, 'P'
        },
        {
            13, 2,
            true, false, 0, 0, F, false, false,
            {1, 1, 0, -1, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            15, 2,
            true, false, -A / 2, -A / 2, 0, false, false,
            {1, 0, 0, 0, 1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            16, 2,
            true, false, D, D, F, true, false,
            {-1, -1, 0, 1, -1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Orthorhombic, 'F'
        },
        {
            14, 2,
            true, false, D, D, F, false, false,
            {1, 1, 0, -1, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            17, 2,
            true, false, D, E, F, true, false,
            {1, -1, 0, 1, 1, 0, -1, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            18, 1,
            false, true, A / 4, A / 2, A / 2, false, false,
            {0, -1, 1, 1, -1, -1, 1, 0, 0},
            gemmi::CrystalSystem::Tetragonal, 'I'
        },
        {
            19, 1,
            false, true, D, A / 2, A / 2, false, false,
            {-1, 0, 0, 0, -1, 1, -1, 1, 1},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            20, 1,
            false, true, D, E, E, false, false,
            {0, 1, 1, 0, 1, -1, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            21, 2,
            false, true, 0, 0, 0, false, false,
            {0, 1, 0, 0, 0, 1, 1, 0, 0},
            gemmi::CrystalSystem::Tetragonal, 'P'
        },
        {
            22, 2,
            false, true, -B / 2, 0, 0, false, false,
            {0, 1, 0, 0, 0, 1, 1, 0, 0},
            gemmi::CrystalSystem::Hexagonal, 'P'
        },
        {
            23, 2,
            false, true, D, 0, 0, false, false,
            {0, 1, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            24, 2,
            false, true, D, -A / 3, -A / 3, true, false,
            {1, 2, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Trigonal, 'R'
        },
        {
            25, 2,
            false, true, D, E, E, false, false,
            {0, 1, 1, 0, -1, 1, 1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            26, 1,
            false, false, A / 4, A / 2, A / 2, false, false,
            {1, 0, 0, -1, 2, 0, -1, 0, 2},
            gemmi::CrystalSystem::Orthorhombic, 'F'
        },
        {
            27, 1,
            false, false, D, A / 2, A / 2, false, false,
            {-1, 2, 0, -1, 0, 0, 0, -1, 1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            28, 1,
            false, false, D, A / 2, 2 * D, false, false,
            {-1, 0, 0, -1, 0, 2, 0, 1, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            29, 1,
            false, false, D, 2 * D, A / 2, false, false,
            {1, 0, 0, 1, -2, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            30, 1,
            false, false, B / 2, E, 2 * E, false, false,
            {0, 1, 0, 0, 1, -2, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },

        {
            31, 1,
            false, false, D, E, F, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Triclinic, 'P'
        },

        {
            32, 2,
            false, false, 0, 0, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Orthorhombic, 'P'
        },
        {
            40, 2,
            false, false, -B / 2, 0, 0, false, false,
            {0, -1, 0, 0, 1, 2, -1, 0, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            35, 2,
            false, false, D, 0, 0, false, false,
            {0, -1, 0, -1, 0, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            36, 2,
            false, false, 0, -A / 2, 0, false, false,
            {1, 0, 0, -1, 0, -2, 0, 1, 0},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            33, 2,
            false, false, 0, E, 0, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            38, 2,
            false, false, 0, 0, -A / 2, false, false,
            {-1, 0, 0, 1, 2, 0, 0, 0, -1},
            gemmi::CrystalSystem::Orthorhombic, 'C'
        },
        {
            34, 2,
            false, false, 0, 0, F, false, false,
            {-1, 0, 0, 0, 0, -1, 0, -1, 0},
            gemmi::CrystalSystem::Monoclinic, 'P'
        },
        {
            42, 2,
            false, false, -B / 2, -A / 2, 0, false, false,
            {-1, 0, 0, 0, -1, 0, 1, 1, 2},
            gemmi::CrystalSystem::Orthorhombic, 'I'
        },
        {
            41, 2,
           false, false, -B / 2, E, 0, false, false,
            {0, -1, -2, 0, -1, 0, -1, 0, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            37, 2,
            false, false, D, -A / 2, 0, false, false,
            {1, 0, 2, 1, 0, 0, 0, 1, 0},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            39, 2,
            false, false, D, 0, -A / 2, false, false,
            {-1, -2, 0, -1, 0, 0, 0, 0, -1},
            gemmi::CrystalSystem::Monoclinic, 'C'
        },
        {
            44, 2,
            false, false, D, E, F, false, false,
            {1, 0, 0, 0, 1, 0, 0, 0, 1},
            gemmi::CrystalSystem::Triclinic, 'P'
        }
      };
    };

    auto match = [&](const CrystalLattice &latt, double D, double E, double F)
            -> std::optional<LatticeSearchResult> {
        const auto uc_reduced = latt.GetUnitCell();
        for (const auto &c: make_classes(D, E, F)) {
            if (only_class && (c.system != only_class->first || c.centering != only_class->second))
                continue;
            if (c.type == 1 && uc_reduced.beta >= 90 - angle_tolerance )
                continue;

            bool ok = true;

            if (c.cond_AB && fabs((uc_reduced.a - uc_reduced.b) / (0.5 * (uc_reduced.a + uc_reduced.b))) > dist_tolerance)
                ok = false;
            if (c.cond_BC && fabs((uc_reduced.b - uc_reduced.c) / (0.5 * (uc_reduced.b + uc_reduced.c))) > dist_tolerance)
                ok = false;

            double expected_alpha = acos(c.cond_D / sqrt(B*C)) * 180 / PI;
            double expected_beta = acos(c.cond_E / sqrt(A*C)) * 180 / PI;
            double expected_gamma = acos(c.cond_F / sqrt(A*B)) * 180 / PI;

            if (fabs(expected_alpha - uc_reduced.alpha) > angle_tolerance)
                ok = false;
            if (fabs(expected_beta - uc_reduced.beta) > angle_tolerance)
                ok = false;
            if (fabs(expected_gamma - uc_reduced.gamma) > angle_tolerance)
                ok = false;

            double tmp1 = 2.0 * fabs(D + E + F);
            double tmp2 = A + B;

            if (c.cond_DEF && fabs((tmp1 - tmp2) / (0.5 * (tmp1 + tmp2))) > dist_tolerance)
                ok = false;

            if (ok) {
                return LatticeSearchResult{
                    .niggli_class = c.number,
                    .primitive_reduced = latt,
                    .conventional = latt.Multiply(c.reindex),
                    .system = c.system,
                    .centering = c.centering,
                        .reindex = c.reindex,
                };
            }
        }
        return std::nullopt;
    };

    // Character 44 fits any cell, so a match is always found - "nothing fits" is reported as triclinic.
    auto found = match(L_niggli, D, E, F);
    if (found && found->system != gemmi::CrystalSystem::Triclinic)
        return *found;

    // A reduced cell whose beta is 90 to within NIGGLI_TYPE_BOUNDARY_DEG sits ON the boundary between the two
    // Niggli types: the same lattice reduces to an all-acute cell or an all-obtuse one according to the
    // last digits of whatever refinement produced it. The type-1 characters are skipped for such a cell
    // (just above) and the type-2 ones are stated for the obtuse setting, so an acute cell can match none
    // of them and comes back triclinic. Present it in the obtuse setting and try once more - negating a
    // and c keeps the lattice and beta and turns alpha and gamma into their supplements. Measured on a
    // C-centred monoclinic crystal whose reduced beta sits 0.07 deg from 90: its centring was read or
    // missed according to the sign of that 0.07 deg, and with it the space group of the whole run.
    if (L_niggli.GetUnitCell().beta >= 90 - NIGGLI_TYPE_BOUNDARY_DEG && D > 0 && E > 0 && F > 0) {
        const gemmi::Mat33 obtuse(-1, 0, 0, 0, 1, 0, 0, 0, -1);
        const auto flipped = match(L_niggli.Multiply(obtuse), -D, E, -F);
        if (flipped && flipped->system != gemmi::CrystalSystem::Triclinic)
            return *flipped;
    }

    if (found)
        return *found;

    if (only_class)
        return std::nullopt;      // no character of the requested class fits this metric

    return LatticeSearchResult{
        .niggli_class = 44,
        .primitive_reduced = L_niggli,
        .conventional = L_niggli,
        .system = gemmi::CrystalSystem::Triclinic,
        .centering = 'P',
        .reindex = gemmi::Mat33(1, 0, 0, 0, 1, 0, 0, 0, 1),
    };
}
}   // namespace

LatticeSearchResult LatticeSearch(const CrystalLattice &L, double dist_tolerance, double angle_tolerance) {
    return *SearchCharacters(L, dist_tolerance, angle_tolerance, nullptr);
}

std::optional<LatticeSearchResult> LatticeSearchForClass(const CrystalLattice &L,
                                                         gemmi::CrystalSystem system, char centering,
                                                         double dist_tolerance, double angle_tolerance) {
    const std::pair<gemmi::CrystalSystem, char> only{system, centering};
    return SearchCharacters(L, dist_tolerance, angle_tolerance, &only);
}