package org.bouncycastle.pqc.crypto.sqisign;

/**
 * Full-rank rational 4-dimensional integer lattice (basis columns scaled by a
 * common denominator). Java port of
 * {@code src/quaternion/ref/generic/lattice.c}.
 *
 * <p>A lattice is represented as a 4x4 integer matrix {@code basis} together
 * with a denominator {@code denom}; the columns of {@code basis / denom} form
 * the algebra elements in the (1, i, j, ij) basis.</p>
 */
final class QuatLattice
{
    public final Ibz denom;
    public final Ibz[][] basis;

    public QuatLattice()
    {
        this.denom = new Ibz(1);
        this.basis = IbzMat.init4x4();
    }

    public QuatLattice copy()
    {
        QuatLattice out = new QuatLattice();
        Ibz.copy(out.denom, this.denom);
        IbzMat.copy4x4(out.basis, this.basis);
        return out;
    }

    /**
     * {@code quat_lattice_reduce_denom}: cancel common factor of basis & denom.
     */
    public static void reduceDenom(QuatLattice reduced, QuatLattice lat)
    {
        Ibz gcd = new Ibz();
        IbzMat.gcd4x4(gcd, lat.basis);
        Ibz.gcd(gcd, gcd, lat.denom);
        IbzMat.scalarDiv4x4(reduced.basis, gcd, lat.basis);
        Ibz.div(reduced.denom, gcd, lat.denom, gcd);
        // Note: the C reuses gcd as the remainder slot for ibz_div; the call
        // overwrites gcd, which is then re-set via abs below. Match that.
        Ibz.abs(reduced.denom, reduced.denom);
    }

    /**
     * {@code quat_lattice_conjugate_without_hnf}: negate rows 1..3.
     */
    public static void conjugateWithoutHnf(QuatLattice conj, QuatLattice lat)
    {
        IbzMat.copy4x4(conj.basis, lat.basis);
        Ibz.copy(conj.denom, lat.denom);
        for (int row = 1; row < 4; row++)
        {
            for (int col = 0; col < 4; col++)
            {
                Ibz.neg(conj.basis[row][col], conj.basis[row][col]);
            }
        }
    }

    /**
     * {@code quat_lattice_dual_without_hnf}: dual lattice via transpose of the
     * scaled inverse. dual.basis = lat.denom * inv(lat.basis); dual.denom = det.
     */
    private static void dualWithoutHnf(QuatLattice dual, QuatLattice lat)
    {
        Ibz det = new Ibz();
        Ibz[][] inv = IbzMat.init4x4();
        IbzMat.invWithDetAsDenom4x4(inv, det, lat.basis);
        IbzMat.transpose4x4(inv, inv);
        IbzMat.scalarMul4x4(dual.basis, lat.denom, inv);
        Ibz.copy(dual.denom, det);
    }

    /**
     * {@code quat_lattice_hnf}: reduce lattice basis to Hermite Normal Form
     * (in place). Mirrors the C reference: invert the basis to obtain the
     * determinant, use |det| as the modulus, build column-generators from the
     * current basis, run HNF, then reduce.
     */
    public static void hnf(QuatLattice lat)
    {
        Ibz mod = new Ibz();
        IbzMat.invWithDetAsDenom4x4(null, mod, lat.basis);
        Ibz.abs(mod, mod);
        Ibz[][] generators = new Ibz[4][4];
        for (int i = 0; i < 4; i++)
        {
            generators[i] = new Ibz[4];
            for (int t = 0; t < 4; t++)
            {
                generators[i][t] = new Ibz();
            }
        }
        for (int i = 0; i < 4; i++)
        {
            for (int j = 0; j < 4; j++)
            {
                Ibz.copy(generators[j][i], lat.basis[i][j]);
            }
        }
        Hnf.hnf4xnModCore(lat.basis, 4, generators, mod);
        reduceDenom(lat, lat);
    }

    /**
     * {@code quat_lattice_equal}: equality after reducing to canonical HNF
     * form with absolute-value denominators.
     */
    public static int equal(QuatLattice lat1, QuatLattice lat2)
    {
        QuatLattice a = new QuatLattice();
        QuatLattice b = new QuatLattice();
        reduceDenom(a, lat1);
        reduceDenom(b, lat2);
        Ibz.abs(a.denom, a.denom);
        Ibz.abs(b.denom, b.denom);
        hnf(a);
        hnf(b);
        if (Ibz.cmp(a.denom, b.denom) != 0)
        {
            return 0;
        }
        return IbzMat.equal4x4(a.basis, b.basis);
    }

    /**
     * {@code quat_lattice_add}: lattice generated by the union of generators.
     * Result is reduced and HNF-normalised.
     */
    public static void add(QuatLattice res, QuatLattice lat1, QuatLattice lat2)
    {
        Ibz[][] tmp = IbzMat.init4x4();
        Ibz[][] generators = new Ibz[8][4];
        for (int h = 0; h < 8; h++)
        {
            generators[h] = new Ibz[4];
            for (int t = 0; t < 4; t++)
            {
                generators[h][t] = new Ibz();
            }
        }
        Ibz det1 = new Ibz();
        Ibz det2 = new Ibz();
        Ibz detprod = new Ibz();

        IbzMat.scalarMul4x4(tmp, lat1.denom, lat2.basis);
        for (int i = 0; i < 4; i++)
        {
            for (int j = 0; j < 4; j++)
            {
                Ibz.copy(generators[j][i], tmp[i][j]);
            }
        }
        IbzMat.invWithDetAsDenom4x4(null, det1, tmp);

        IbzMat.scalarMul4x4(tmp, lat2.denom, lat1.basis);
        for (int i = 0; i < 4; i++)
        {
            for (int j = 0; j < 4; j++)
            {
                Ibz.copy(generators[4 + j][i], tmp[i][j]);
            }
        }
        IbzMat.invWithDetAsDenom4x4(null, det2, tmp);

        if (Ibz.isZero(det1) == 1 || Ibz.isZero(det2) == 1)
        {
            throw new IllegalStateException("lattice add: singular basis");
        }

        Ibz.gcd(detprod, det1, det2);
        Hnf.hnf4xnModCore(res.basis, 8, generators, detprod);
        Ibz.mul(res.denom, lat1.denom, lat2.denom);
        reduceDenom(res, res);
    }

    /**
     * {@code quat_lattice_intersect}: intersection via dual-add-dual.
     * Returns an HNF-normalised result.
     */
    public static void intersect(QuatLattice res, QuatLattice lat1, QuatLattice lat2)
    {
        QuatLattice d1 = new QuatLattice();
        QuatLattice d2 = new QuatLattice();
        QuatLattice dRes = new QuatLattice();
        dualWithoutHnf(d1, lat1);
        dualWithoutHnf(d2, lat2);
        add(dRes, d1, d2);
        dualWithoutHnf(res, dRes);
        hnf(res);
    }

    /**
     * {@code quat_lattice_contains}: returns 1 iff the quaternion algebra
     * element {@code x} lies in {@code lat}. When the return is 1 and
     * {@code coord != null}, the integer coordinates of {@code x} in the
     * lattice's basis are written there.
     */
    public static int contains(Ibz[] coord, QuatLattice lat, QuatAlg.Elem x)
    {
        Ibz det = new Ibz();
        Ibz prod = new Ibz();
        Ibz[] workCoord = IbzVec.init4();
        Ibz[][] inv = IbzMat.init4x4();
        IbzMat.invWithDetAsDenom4x4(inv, det, lat.basis);
        if (Ibz.isZero(det) == 1)
        {
            throw new IllegalStateException("contains: singular basis");
        }
        IbzMat.eval4x4(workCoord, inv, x.coord);
        IbzVec.scalarMul4(workCoord, lat.denom, workCoord);
        Ibz.mul(prod, x.denom, det);
        Ibz[] quotient = IbzVec.init4();
        int divisible = IbzVec.scalarDiv4(quotient, prod, workCoord);
        if (divisible == 1 && coord != null)
        {
            for (int i = 0; i < 4; i++)
            {
                Ibz.copy(coord[i], quotient[i]);
            }
        }
        return divisible;
    }

    /**
     * {@code quat_lattice_index}: index of sublat in overlat (a positive
     * integer when sublat ⊆ overlat).
     */
    public static void index(Ibz index, QuatLattice sublat, QuatLattice overlat)
    {
        Ibz tmp = new Ibz();
        Ibz det = new Ibz();

        IbzMat.invWithDetAsDenom4x4(null, det, sublat.basis);
        Ibz.mul(tmp, overlat.denom, overlat.denom);
        Ibz.mul(tmp, tmp, tmp);
        Ibz.mul(index, det, tmp);

        Ibz.mul(tmp, sublat.denom, sublat.denom);
        Ibz.mul(tmp, tmp, tmp);
        IbzMat.invWithDetAsDenom4x4(null, det, overlat.basis);
        Ibz.mul(tmp, tmp, det);
        Ibz.div(index, tmp, index, tmp);
        Ibz.abs(index, index);
    }

    /**
     * Multiply each column of the basis (treated as an algebra element in
     * basis (1,i,j,ij)) by an algebra-coordinate vector {@code coord} on the
     * right. Mirrors {@code quat_lattice_mat_alg_coord_mul_without_hnf}.
     */
    private static void matAlgCoordMulWithoutHnf(Ibz[][] prod, Ibz[][] lat, Ibz[] coord, QuatAlg alg)
    {
        Ibz[] p = IbzVec.init4();
        Ibz[] a = IbzVec.init4();
        for (int i = 0; i < 4; i++)
        {
            // a = column i of lat
            for (int t = 0; t < 4; t++)
            {
                Ibz.copy(a[t], lat[t][i]);
            }
            QuatAlg.coordMul(p, a, coord, alg);
            for (int t = 0; t < 4; t++)
            {
                Ibz.copy(prod[t][i], p[t]);
            }
        }
    }

    /**
     * {@code quat_lattice_alg_elem_mul}: multiply the lattice by an algebra
     * element on the right. The product is HNF-reduced.
     */
    public static void algElemMul(QuatLattice prod, QuatLattice lat, QuatAlg.Elem elem, QuatAlg alg)
    {
        matAlgCoordMulWithoutHnf(prod.basis, lat.basis, elem.coord, alg);
        Ibz.mul(prod.denom, lat.denom, elem.denom);
        hnf(prod);
    }

    /**
     * {@code quat_lattice_mul}: product of two lattices in the algebra.
     * Result is HNF-normalised.
     */
    public static void mul(QuatLattice res, QuatLattice lat1, QuatLattice lat2, QuatAlg alg)
    {
        Ibz[] elem1 = IbzVec.init4();
        Ibz[] elem2 = IbzVec.init4();
        Ibz[] elemRes = IbzVec.init4();
        Ibz[][] generators = new Ibz[16][4];
        for (int h = 0; h < 16; h++)
        {
            generators[h] = new Ibz[4];
            for (int t = 0; t < 4; t++)
            {
                generators[h][t] = new Ibz();
            }
        }
        Ibz[][] detmat = IbzMat.init4x4();
        Ibz det = new Ibz();

        for (int k = 0; k < 4; k++)
        {
            for (int t = 0; t < 4; t++)
            {
                Ibz.copy(elem1[t], lat1.basis[t][k]);
            }
            for (int i = 0; i < 4; i++)
            {
                for (int t = 0; t < 4; t++)
                {
                    Ibz.copy(elem2[t], lat2.basis[t][i]);
                }
                QuatAlg.coordMul(elemRes, elem1, elem2, alg);
                for (int j = 0; j < 4; j++)
                {
                    if (k == 0)
                    {
                        Ibz.copy(detmat[i][j], elemRes[j]);
                    }
                    Ibz.copy(generators[4 * k + i][j], elemRes[j]);
                }
            }
        }
        IbzMat.invWithDetAsDenom4x4(null, det, detmat);
        Ibz.abs(det, det);
        Hnf.hnf4xnModCore(res.basis, 16, generators, det);
        Ibz.mul(res.denom, lat1.denom, lat2.denom);
        reduceDenom(res, res);
    }

    /**
     * {@code quat_lattice_gram}: Gram matrix of the lattice basis with respect
     * to the quaternion norm form on {(1, i, j, ij)}: G[i][j] = 2 * trace(b_i * conj(b_j)).
     */
    public static void gram(Ibz[][] G, QuatLattice lattice, QuatAlg alg)
    {
        Ibz tmp = new Ibz();
        Ibz two = new Ibz(2);
        for (int i = 0; i < 4; i++)
        {
            for (int j = 0; j <= i; j++)
            {
                Ibz.set(G[i][j], 0);
                for (int k = 0; k < 4; k++)
                {
                    Ibz.mul(tmp, lattice.basis[k][i], lattice.basis[k][j]);
                    if (k >= 2)
                    {
                        Ibz.mul(tmp, tmp, alg.p);
                    }
                    Ibz.add(G[i][j], G[i][j], tmp);
                }
                Ibz.mul(G[i][j], G[i][j], two);
            }
        }
        // Reflect into upper triangle
        for (int i = 0; i < 4; i++)
        {
            for (int j = i + 1; j < 4; j++)
            {
                Ibz.copy(G[i][j], G[j][i]);
            }
        }
    }
}
