"""
A small SageMath script to generate the macro parameters for a given prime
"""

from sage.all import GF, proof, ceil

proof.all(False)


def fmt_little_u64(res, words):
    out = []
    for r in res:
        num = r[2:]
        new_num = num.zfill(16)
        new_num = new_num.upper()
        new_num = "0x" + new_num
        out.append(new_num)
    while len(out) != words:
        out.append("0x" + "0" * 16)
    return out


def fmt_list(l):
    l = str(l)
    l = l.replace("[", "")
    l = l.replace("]", "")
    l = l.replace("'", "")
    return l


def to_little_u64(n, words):

    res = []
    while n:
        tmp = n % 2**64
        res.append(hex(tmp))
        n >>= 64
    l = fmt_little_u64(res, words)
    return fmt_list(l)


if __name__ == "__main__":
    # p = 65 * 2**376 - 1
    # p = 2**192 * 3**243 * 5**28 * 7**2 - 1
    # p = 3**2 * 7 * 2**4084 - 1
    p =826791736418446924644415105270960270928927659729776400179861442336062222833458285859
    assert p.is_prime()
    assert p % 4 == 3
    words = ceil(p.nbits() / 64)
    BITLEN = p.nbits()
    MODULUS = to_little_u64(p, words)

    Fp2 = GF(p**2, names="i", modulus=[1, 0, 1])
    nqr = 1
    for i in range(1, 1000):
        x = Fp2([i, 1])
        if not x.is_square():
            nqr = i
            break

    output = f"""
mod fp{BITLEN}_tests {{
    // Field modulus
    static MODULUS: [u64; {words}] = [{MODULUS}];

    fp2::define_fp_core!(
        typename = Fp{BITLEN},
        modulus = MODULUS,
    );

    fp2::define_fp2_from_type!(
        typename = Fp{BITLEN}Ext,
        base_field = Fp{BITLEN},
    );

    fp2::define_fp_tests!(Fp{BITLEN});
    fp2::define_fp2_tests!(Fp{BITLEN}Ext, MODULUS, {nqr});
}}
    """

    print(output)
