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

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

#include <cmath>

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

LatticeSearchResult LatticeSearch(const CrystalLattice &L, double dist_tolerance, double angle_tolerance) {
    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;

    std::vector<NiggliClass> niggli_classes = {
        {
            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, -1, 0, 0, 0, 0, -1},
            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'
        }
    };

    const auto uc_reduced = L_niggli.GetUnitCell();

    for (const auto &c: niggli_classes) {
        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 / M_PI;
        double expected_beta = acos(c.cond_E / sqrt(A*C)) * 180 / M_PI;
        double expected_gamma = acos(c.cond_F / sqrt(A*B)) * 180 / M_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 = L,
                .conventional = L_niggli.Multiply(c.reindex),
                .system = c.system,
                .centering = c.centering,
            };
        }
    }

    return LatticeSearchResult{
        .niggli_class = 44,
        .primitive_reduced = L,
        .conventional = L,
        .system = gemmi::CrystalSystem::Triclinic,
        .centering = 'P',
    };
}