"""
Step 4: WR_CRC -- DDR4/5 write-CRC integrity.

CRC is linear over GF(2), so a full-burst CRC over 64+ bits can never be
enumerated -- but a single-bit slice can. We verify the CRC BIT-SLICE
exhaustively (state 8b + data bit -> new state, 512 configs) and ripple it,
exactly the byte-slice/ripple approach of the neural-aarch64 datapath.

Polynomial is configurable; default 0x07. Swap it to match the target JEDEC spec.
"""
from __future__ import annotations
import torch
from .common import bits_of, int_of, pm

POLY = 0x07
IN, OUT = 9, 8


def crc_bit(state: int, databit: int) -> int:
    fb = ((state >> 7) & 1) ^ (databit & 1)
    ns = (state << 1) & 0xFF
    if fb:
        ns ^= POLY
    return ns & 0xFF


def domain():
    xs, ys = [], []
    for s in range(256):
        for d in (0, 1):
            v = (d << 8) | s                       # bits0-7=state, bit8=databit
            xs.append(pm(bits_of(v, 9)))
            ys.append([float(b) for b in bits_of(crc_bit(s, d), 8)])
    return torch.stack(xs), torch.tensor(ys)


def run_bit(net, state: int, databit: int) -> int:
    v = (databit << 8) | state
    return int_of((net(pm(bits_of(v, 9)).unsqueeze(0))[0] > 0).int().tolist())


def crc_byte(net, state: int, byte: int) -> int:
    for k in range(7, -1, -1):                     # MSB-first
        state = run_bit(net, state, (byte >> k) & 1)
    return state


def crc_burst(net, data_bytes, init=0) -> int:
    s = init
    for b in data_bytes:
        s = crc_byte(net, s, b)
    return s


def golden_burst(data_bytes, init=0) -> int:
    s = init
    for b in data_bytes:
        for k in range(7, -1, -1):
            s = crc_bit(s, (b >> k) & 1)
    return s