"""
Step 1 keystone: verified GF(2^8) multiply -- the arithmetic core of Reed-Solomon
erasure coding (Step 2).

A direct 16-bit -> 8-bit neural multiply is hard to drive to N/N. So we decompose
it the way real RS implementations do -- via log/exp tables:

    a * b = EXP[ (LOG[a] + LOG[b]) mod 255 ]      (and 0 if a==0 or b==0)

LOG and EXP are 256-entry lookups -> trivially N/N verifiable neural units. We
train both, compose them into a full multiply, and verify the COMPOSED multiply
against the golden field over ALL 65,536 (a,b) pairs.

Field: GF(2^8), reduction polynomial 0x11D, generator 2 (standard for RS).
"""
from __future__ import annotations
import torch
from .common import mlp, train, verify, pm, bits_of, int_of, run8

POLY = 0x11D


def build_tables():
    exp = [0] * 512
    log = [0] * 256
    x = 1
    for i in range(255):
        exp[i] = x
        log[x] = i
        x <<= 1
        if x & 0x100:
            x ^= POLY
    for i in range(255, 512):
        exp[i] = exp[i - 255]
    return exp, log


EXP, LOG = build_tables()


def gf_mul(a: int, b: int) -> int:
    if a == 0 or b == 0:
        return 0
    return EXP[LOG[a] + LOG[b]]


# ---- neural LOG (inputs 1..255) and EXP (indices 0..254) ----
def log_domain():
    xs = [pm(bits_of(a, 8)) for a in range(1, 256)]
    ys = [[float(b) for b in bits_of(LOG[a], 8)] for a in range(1, 256)]
    return torch.stack(xs), torch.tensor(ys)


def exp_domain():
    xs = [pm(bits_of(i, 8)) for i in range(255)]
    ys = [[float(b) for b in bits_of(EXP[i], 8)] for i in range(255)]
    return torch.stack(xs), torch.tensor(ys)


def gf_mul_neural(a: int, b: int, net_log, net_exp) -> int:
    if a == 0 or b == 0:
        return 0
    s = (run8(net_log, a) + run8(net_log, b)) % 255
    return run8(net_exp, s)


def train_units():
    print("  training LOG (8->8):")
    net_log = train(mlp(8, 8, 256, 2), *log_domain(), steps=8000, tag="log")
    print("  training EXP (8->8):")
    net_exp = train(mlp(8, 8, 256, 2), *exp_domain(), steps=8000, tag="exp")
    return net_log, net_exp


def verify_mul(net_log, net_exp) -> tuple[int, int]:
    ok = 0
    for a in range(256):
        for b in range(256):
            ok += (gf_mul_neural(a, b, net_log, net_exp) == gf_mul(a, b))
    return ok, 256 * 256