// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Ported to V from Go's crypto/internal/fips140/mldsa.
module mldsa

// appendix b: zeta^BitRev8(k) mod q in montgomery domain
const zetas = [u32(4193792), 25847, 5771523, 7861508, 237124, 7602457, 7504169, 466468, 1826347,
    2353451, 8021166, 6288512, 3119733, 5495562, 3111497, 2680103, 2725464, 1024112, 7300517, 3585928,
    7830929, 7260833, 2619752, 6271868, 6262231, 4520680, 6980856, 5102745, 1757237, 8360995, 4010497,
    280005, 2706023, 95776, 3077325, 3530437, 6718724, 4788269, 5842901, 3915439, 4519302, 5336701,
    3574422, 5512770, 3539968, 8079950, 2348700, 7841118, 6681150, 6736599, 3505694, 4558682, 3507263,
    6239768, 6779997, 3699596, 811944, 531354, 954230, 3881043, 3900724, 5823537, 2071892, 5582638,
    4450022, 6851714, 4702672, 5339162, 6927966, 3475950, 2176455, 6795196, 7122806, 1939314, 4296819,
    7380215, 5190273, 5223087, 4747489, 126922, 3412210, 7396998, 2147896, 2715295, 5412772, 4686924,
    7969390, 5903370, 7709315, 7151892, 8357436, 7072248, 7998430, 1349076, 1852771, 6949987, 5037034,
    264944, 508951, 3097992, 44288, 7280319, 904516, 3958618, 4656075, 8371839, 1653064, 5130689,
    2389356, 8169440, 759969, 7063561, 189548, 4827145, 3159746, 6529015, 5971092, 8202977, 1315589,
    1341330, 1285669, 6795489, 7567685, 6940675, 5361315, 4499357, 4751448, 3839961, 2091667, 3407706,
    2316500, 3817976, 5037939, 2244091, 5933984, 4817955, 266997, 2434439, 7144689, 3513181, 4860065,
    4621053, 7183191, 5187039, 900702, 1859098, 909542, 819034, 495491, 6767243, 8337157, 7857917,
    7725090, 5257975, 2031748, 3207046, 4823422, 7855319, 7611795, 4784579, 342297, 286988, 5942594,
    4108315, 3437287, 5038140, 1735879, 203044, 2842341, 2691481, 5790267, 1265009, 4055324, 1247620,
    2486353, 1595974, 4613401, 1250494, 2635921, 4832145, 5386378, 1869119, 1903435, 7329447, 7047359,
    1237275, 5062207, 6950192, 7929317, 1312455, 3306115, 6417775, 7100756, 1917081, 5834105, 7005614,
    1500165, 777191, 2235880, 3406031, 7838005, 5548557, 6709241, 6533464, 5796124, 4656147, 594136,
    4603424, 6366809, 2432395, 2454455, 8215696, 1957272, 3369112, 185531, 7173032, 5196991, 162844,
    1616392, 3014001, 810149, 1652634, 4686184, 6581310, 5341501, 3523897, 3866901, 269760, 2213111,
    7404533, 1717735, 472078, 7953734, 1723600, 6577327, 1910376, 6712985, 7276084, 8119771, 4546524,
    5441381, 6144432, 7959518, 6094090, 183443, 7403526, 1612842, 4834730, 7826001, 3919660, 8332111,
    7018208, 3937738, 1400424, 7534263, 1976782]!

// algo. 41: NTT (s. 7.5)
@[direct_array_access]
fn ntt(f_ RingElement) NttElement {
    mut f := unsafe { f_ } // workaround for false mutability notice
    mut m := u8(0)

    mut len := 128
    for len >= 8 {
        mut start := 0
        for start < 256 {
            m++
            zeta := FieldElement(zetas[m])

            mut j := 0
            for j < len {
                t := field_montgomery_mul(zeta, f[start + len + j])
                f[start + len + j] = field_sub(f[start + j], t)
                f[start + j] = field_add(f[start + j], t)

                t2 := field_montgomery_mul(zeta, f[start + len + j + 1])
                f[start + len + j + 1] = field_sub(f[start + j + 1], t2)
                f[start + j + 1] = field_add(f[start + j + 1], t2)

                j += 2
            }
            start += 2 * len
        }
        len /= 2
    }

    for start := 0; start < 256; start += 8 {
        m++
        zeta := FieldElement(zetas[m])

        mut t := field_montgomery_mul(zeta, f[start + 4])
        f[start + 4] = field_sub(f[start], t)
        f[start] = field_add(f[start], t)

        t = field_montgomery_mul(zeta, f[start + 5])
        f[start + 5] = field_sub(f[start + 1], t)
        f[start + 1] = field_add(f[start + 1], t)

        t = field_montgomery_mul(zeta, f[start + 6])
        f[start + 6] = field_sub(f[start + 2], t)
        f[start + 2] = field_add(f[start + 2], t)

        t = field_montgomery_mul(zeta, f[start + 7])
        f[start + 7] = field_sub(f[start + 3], t)
        f[start + 3] = field_add(f[start + 3], t)
    }

    for start := 0; start < 256; start += 4 {
        m++
        zeta := FieldElement(zetas[m])

        mut t := field_montgomery_mul(zeta, f[start + 2])
        f[start + 2] = field_sub(f[start], t)
        f[start] = field_add(f[start], t)

        t = field_montgomery_mul(zeta, f[start + 3])
        f[start + 3] = field_sub(f[start + 1], t)
        f[start + 1] = field_add(f[start + 1], t)
    }

    for start := 0; start < 256; start += 2 {
        m++
        zeta := FieldElement(zetas[m])

        t := field_montgomery_mul(zeta, f[start + 1])
        f[start + 1] = field_sub(f[start], t)
        f[start] = field_add(f[start], t)
    }

    return NttElement(f)
}

// algo. 42: NTT^-1 (s. 7.5)
@[direct_array_access]
fn inverse_ntt(f_ NttElement) RingElement {
    mut f := unsafe { f_ } // workaround for false mutability notice
    mut m := u8(255)

    for start := 0; start < 256; start += 2 {
        zeta := FieldElement(zetas[m])
        m--

        t := f[start]
        f[start] = field_add(t, f[start + 1])
        f[start + 1] = field_montgomery_mul_sub(zeta, f[start + 1], t)
    }

    for start := 0; start < 256; start += 4 {
        zeta := FieldElement(zetas[m])
        m--

        mut t := f[start]
        f[start] = field_add(t, f[start + 2])
        f[start + 2] = field_montgomery_mul_sub(zeta, f[start + 2], t)

        t = f[start + 1]
        f[start + 1] = field_add(t, f[start + 3])
        f[start + 3] = field_montgomery_mul_sub(zeta, f[start + 3], t)
    }

    for start := 0; start < 256; start += 8 {
        zeta := FieldElement(zetas[m])
        m--

        mut t := f[start]
        f[start] = field_add(t, f[start + 4])
        f[start + 4] = field_montgomery_mul_sub(zeta, f[start + 4], t)

        t = f[start + 1]
        f[start + 1] = field_add(t, f[start + 5])
        f[start + 5] = field_montgomery_mul_sub(zeta, f[start + 5], t)

        t = f[start + 2]
        f[start + 2] = field_add(t, f[start + 6])
        f[start + 6] = field_montgomery_mul_sub(zeta, f[start + 6], t)

        t = f[start + 3]
        f[start + 3] = field_add(t, f[start + 7])
        f[start + 7] = field_montgomery_mul_sub(zeta, f[start + 7], t)
    }

    mut len2 := 8
    for len2 < 256 {
        mut start := 0
        for start < 256 {
            zeta := FieldElement(zetas[m])
            m--

            mut j := 0
            for j < len2 {
                mut t := f[start + j]
                f[start + j] = field_add(t, f[start + len2 + j])
                f[start + len2 + j] = field_montgomery_mul_sub(zeta, f[start + len2 + j], t)

                t = f[start + j + 1]
                f[start + j + 1] = field_add(t, f[start + len2 + j + 1])
                f[start + len2 + j + 1] = field_montgomery_mul_sub(zeta, f[start + len2 + j + 1], t)

                j += 2
            }
            start += 2 * len2
        }
        len2 *= 2
    }

    // algo. 42, line 21: f = 8347681 = 256^-1 mod q; in montgomery: 16382
    for i in 0 .. 256 {
        f[i] = field_montgomery_mul(f[i], 16382)
    }
    return RingElement(f)
}