#include <quaternion.h>
#include <stdlib.h>
#include "internal.h"

// assumes parent order and lattice correctly set, computes and sets the norm
void
quat_lideal_norm(quat_left_ideal_t *lideal)
{
    quat_lattice_index(&(lideal->norm), &(lideal->lattice), (lideal->parent_order));
    int ok UNUSED = ibz_sqrt(&(lideal->norm), &(lideal->norm));
    assert(ok);
}

// assumes parent order and lattice correctly set, recomputes and verifies its norm
static int
quat_lideal_norm_verify(const quat_left_ideal_t *lideal)
{
    int res;
    ibz_t index;
    ibz_init(&index);
    quat_lattice_index(&index, &(lideal->lattice), (lideal->parent_order));
    ibz_sqrt(&index, &index);
    res = (ibz_cmp(&(lideal->norm), &index) == 0);
    ibz_finalize(&index);
    return (res);
}

void
quat_lideal_copy(quat_left_ideal_t *copy, const quat_left_ideal_t *copied)
{
    copy->parent_order = copied->parent_order;
    ibz_copy(&copy->norm, &copied->norm);
    ibz_copy(&copy->lattice.denom, &copied->lattice.denom);
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            ibz_copy(&copy->lattice.basis[i][j], &copied->lattice.basis[i][j]);
        }
    }
}

void
quat_lideal_create_principal(quat_left_ideal_t *lideal,
                             const quat_alg_elem_t *x,
                             const quat_lattice_t *order,
                             const quat_alg_t *alg)
{
    assert(quat_order_is_maximal(order, alg));
    assert(quat_lattice_contains(NULL, order, x));
    ibz_t norm_n, norm_d;
    ibz_init(&norm_n);
    ibz_init(&norm_d);

    // Multiply order on the right by x
    quat_lattice_alg_elem_mul(&(lideal->lattice), order, x, alg);

    // Reduce denominator. This conserves HNF
    quat_lattice_reduce_denom(&lideal->lattice, &lideal->lattice);

    // Compute norm and check it's integral
    quat_alg_norm(&norm_n, &norm_d, x, alg);
    assert(ibz_is_one(&norm_d));
    ibz_copy(&lideal->norm, &norm_n);

    // Set order
    lideal->parent_order = order;
    ibz_finalize(&norm_n);
    ibz_finalize(&norm_d);
}

void
quat_lideal_create(quat_left_ideal_t *lideal,
                   const quat_alg_elem_t *x,
                   const ibz_t *N,
                   const quat_lattice_t *order,
                   const quat_alg_t *alg)
{
    assert(quat_order_is_maximal(order, alg));
    assert(!quat_alg_elem_is_zero(x));

    quat_lattice_t ON;
    quat_lattice_init(&ON);

    // Compute ideal generated by x
    quat_lideal_create_principal(lideal, x, order, alg);

    // Compute ideal generated by N (without reducing denominator)
    ibz_mat_4x4_scalar_mul(&ON.basis, N, &order->basis);
    ibz_copy(&ON.denom, &order->denom);

    // Add lattices (reduces denominators)
    quat_lattice_add(&lideal->lattice, &lideal->lattice, &ON);
    // Set order
    lideal->parent_order = order;
    // Compute norm
    quat_lideal_norm(lideal);

    quat_lattice_finalize(&ON);
}

int
quat_lideal_generator(quat_alg_elem_t *gen, const quat_left_ideal_t *lideal, const quat_alg_t *alg)
{
    ibz_t norm_int, norm_n, gcd, r, q, norm_denom;
    ibz_vec_4_t vec;
    ibz_vec_4_init(&vec);
    ibz_init(&norm_denom);
    ibz_init(&norm_int);
    ibz_init(&norm_n);
    ibz_init(&r);
    ibz_init(&q);
    ibz_init(&gcd);
    int a, b, c, d;
    int found = 0;
    int int_norm = 0;
    while (1) {
        int_norm++;
        for (a = -int_norm; a <= int_norm; a++) {
            for (b = -int_norm + abs(a); b <= int_norm - abs(a); b++) {
                for (c = -int_norm + abs(a) + abs(b); c <= int_norm - abs(a) - abs(b); c++) {
                    d = int_norm - abs(a) - abs(b) - abs(c);
                    ibz_vec_4_set(&vec, a, b, c, d);
                    ibz_vec_4_content(&gcd, &vec);
                    if (ibz_is_one(&gcd)) {
                        ibz_mat_4x4_eval(&(gen->coord), &(lideal->lattice.basis), &vec);
                        ibz_copy(&(gen->denom), &(lideal->lattice.denom));
                        quat_alg_norm(&norm_int, &norm_denom, gen, alg);
                        assert(ibz_is_one(&norm_denom));
                        ibz_div(&q, &r, &norm_int, &(lideal->norm));
                        assert(ibz_is_zero(&r));
                        ibz_gcd(&gcd, &(lideal->norm), &q);
                        found = (0 == ibz_cmp(&gcd, &ibz_const_one));
                        if (found)
                            goto fin;
                    }
                }
            }
        }
    }
fin:;
    ibz_finalize(&r);
    ibz_finalize(&q);
    ibz_finalize(&norm_denom);
    ibz_finalize(&norm_int);
    ibz_finalize(&norm_n);
    ibz_vec_4_finalize(&vec);
    ibz_finalize(&gcd);
    return (found);
}

void
quat_lideal_mul(quat_left_ideal_t *product,
                const quat_left_ideal_t *lideal,
                const quat_alg_elem_t *alpha,
                const quat_alg_t *alg)
{
    assert(quat_order_is_maximal((lideal->parent_order), alg));
    ibz_t norm, norm_d;
    ibz_init(&norm);
    ibz_init(&norm_d);
    quat_lattice_alg_elem_mul(&(product->lattice), &(lideal->lattice), alpha, alg);
    product->parent_order = lideal->parent_order;
    quat_alg_norm(&norm, &norm_d, alpha, alg);
    ibz_mul(&(product->norm), &(lideal->norm), &norm);
    assert(ibz_divides(&(product->norm), &norm_d));
    ibz_div(&(product->norm), &norm, &(product->norm), &norm_d);
    assert(quat_lideal_norm_verify(lideal));
    ibz_finalize(&norm_d);
    ibz_finalize(&norm);
}

void
quat_lideal_add(quat_left_ideal_t *sum, const quat_left_ideal_t *I1, const quat_left_ideal_t *I2, const quat_alg_t *alg)
{
    assert(I1->parent_order == I2->parent_order);
    assert(quat_order_is_maximal((I2->parent_order), alg));
    quat_lattice_add(&sum->lattice, &I1->lattice, &I2->lattice);
    sum->parent_order = I1->parent_order;
    quat_lideal_norm(sum);
}

void
quat_lideal_inter(quat_left_ideal_t *inter,
                  const quat_left_ideal_t *I1,
                  const quat_left_ideal_t *I2,
                  const quat_alg_t *alg)
{
    assert(I1->parent_order == I2->parent_order);
    assert(quat_order_is_maximal((I2->parent_order), alg));
    quat_lattice_intersect(&inter->lattice, &I1->lattice, &I2->lattice);
    inter->parent_order = I1->parent_order;
    quat_lideal_norm(inter);
}

int
quat_lideal_equals(const quat_left_ideal_t *I1, const quat_left_ideal_t *I2, const quat_alg_t *alg)
{
    assert(quat_order_is_maximal((I2->parent_order), alg));
    assert(quat_order_is_maximal((I1->parent_order), alg));
    return (I1->parent_order == I2->parent_order) & (ibz_cmp(&I1->norm, &I2->norm) == 0) &
           quat_lattice_equal(&I1->lattice, &I2->lattice);
}

void
quat_lideal_inverse_lattice_without_hnf(quat_lattice_t *inv, const quat_left_ideal_t *lideal, const quat_alg_t *alg)
{
    assert(quat_order_is_maximal((lideal->parent_order), alg));
    quat_lattice_conjugate_without_hnf(inv, &(lideal->lattice));
    ibz_mul(&(inv->denom), &(inv->denom), &(lideal->norm));
}

// following the implementation of ideal isomorphisms in the code of LearningToSQI's sage
// implementation of SQIsign
void
quat_lideal_right_transporter(quat_lattice_t *trans,
                              const quat_left_ideal_t *lideal1,
                              const quat_left_ideal_t *lideal2,
                              const quat_alg_t *alg)
{
    assert(quat_order_is_maximal((lideal1->parent_order), alg));
    assert(quat_order_is_maximal((lideal2->parent_order), alg));
    assert(lideal1->parent_order == lideal2->parent_order);
    quat_lattice_t inv;
    quat_lattice_init(&inv);
    quat_lideal_inverse_lattice_without_hnf(&inv, lideal1, alg);
    quat_lattice_mul(trans, &inv, &(lideal2->lattice), alg);
    quat_lattice_finalize(&inv);
}

void
quat_lideal_right_order(quat_lattice_t *order, const quat_left_ideal_t *lideal, const quat_alg_t *alg)
{
    assert(quat_order_is_maximal((lideal->parent_order), alg));
    quat_lideal_right_transporter(order, lideal, lideal, alg);
}

void
quat_lideal_class_gram(ibz_mat_4x4_t *G, const quat_left_ideal_t *lideal, const quat_alg_t *alg)
{
    quat_lattice_gram(G, &(lideal->lattice), alg);

    // divide by norm · denominator²
    ibz_t divisor, rmd;
    ibz_init(&divisor);
    ibz_init(&rmd);

    ibz_mul(&divisor, &(lideal->lattice.denom), &(lideal->lattice.denom));
    ibz_mul(&divisor, &divisor, &(lideal->norm));

    for (int i = 0; i < 4; i++) {
        for (int j = 0; j <= i; j++) {
            ibz_div(&(*G)[i][j], &rmd, &(*G)[i][j], &divisor);
            assert(ibz_is_zero(&rmd));
        }
    }
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j <= i - 1; j++) {
            ibz_copy(&(*G)[j][i], &(*G)[i][j]);
        }
    }

    ibz_finalize(&rmd);
    ibz_finalize(&divisor);
}

void
quat_lideal_conjugate_without_hnf(quat_left_ideal_t *conj,
                                  quat_lattice_t *new_parent_order,
                                  const quat_left_ideal_t *lideal,
                                  const quat_alg_t *alg)
{
    quat_lideal_right_order(new_parent_order, lideal, alg);
    quat_lattice_conjugate_without_hnf(&(conj->lattice), &(lideal->lattice));
    conj->parent_order = new_parent_order;
    ibz_copy(&(conj->norm), &(lideal->norm));
}

int
quat_order_discriminant(ibz_t *disc, const quat_lattice_t *order, const quat_alg_t *alg)
{
    int ok = 0;
    ibz_t det, sqr, div;
    ibz_mat_4x4_t transposed, norm, prod;
    ibz_init(&det);
    ibz_init(&sqr);
    ibz_init(&div);
    ibz_mat_4x4_init(&transposed);
    ibz_mat_4x4_init(&norm);
    ibz_mat_4x4_init(&prod);
    ibz_mat_4x4_transpose(&transposed, &(order->basis));
    // multiply gram matrix by 2 because of reduced trace
    ibz_mat_4x4_identity(&norm);
    ibz_copy(&(norm[2][2]), &(alg->p));
    ibz_copy(&(norm[3][3]), &(alg->p));
    ibz_mat_4x4_scalar_mul(&norm, &ibz_const_two, &norm);
    ibz_mat_4x4_mul(&prod, &transposed, &norm);
    ibz_mat_4x4_mul(&prod, &prod, &(order->basis));
    ibz_mat_4x4_inv_with_det_as_denom(NULL, &det, &prod);
    ibz_mul(&div, &(order->denom), &(order->denom));
    ibz_mul(&div, &div, &div);
    ibz_mul(&div, &div, &div);
    ibz_div(&sqr, &div, &det, &div);
    ok = ibz_is_zero(&div);
    ok = ok & ibz_sqrt(disc, &sqr);
    ibz_finalize(&det);
    ibz_finalize(&div);
    ibz_finalize(&sqr);
    ibz_mat_4x4_finalize(&transposed);
    ibz_mat_4x4_finalize(&norm);
    ibz_mat_4x4_finalize(&prod);
    return (ok);
}

int
quat_order_is_maximal(const quat_lattice_t *order, const quat_alg_t *alg)
{
    int res;
    ibz_t disc;
    ibz_init(&disc);
    quat_order_discriminant(&disc, order, alg);
    res = (ibz_cmp(&disc, &(alg->p)) == 0);
    ibz_finalize(&disc);
    return (res);
}
