"""
Step 1 of neural-DDR: the DBI (Data Bus Inversion) unit, verified N/N.

Same methodology as the neural-aarch64 datapath units: DBI is a FIXED finite
function, so a small MLP can be trained until it is bit-identical to a golden
reference over its ENTIRE input domain (exhaustive N/N verification).

DBI is a real DDR4/DDR5 feature (absent in DDR3) -- so it is exactly the kind of
per-generation logic the bridge swaps in/out. On the data bus, DDR reads back
data over DQ lines that are terminated to a pull-up; a line driven LOW burns
termination power. DBI_dc reduces that: if a byte has more than four 0-bits,
invert all eight and flag it, so at most four lines are ever LOW.

  encode(byte) -> 9 bits = [8 encoded data][1 DBI flag]
      inverted = (number_of_zero_bits(byte) > 4)
      data_out = ~byte & 0xFF  if inverted else byte    # <=4 zeros afterwards
      flag     = 1 if inverted else 0
  decode(data_out, flag) -> byte
      byte = ~data_out & 0xFF  if flag else data_out

Neither storage nor capacity is invented here: this unit only emulates the
bus-encoding LOGIC. The bytes themselves live in the host's real RAM (bridge).
"""

from __future__ import annotations

import torch
import torch.nn as nn


# ---------- golden reference (the finite function we must match) ----------
def golden_encode(byte: int) -> tuple[int, int]:
    zeros = 8 - bin(byte & 0xFF).count("1")
    if zeros > 4:
        return (~byte) & 0xFF, 1
    return byte & 0xFF, 0


def golden_decode(data_out: int, flag: int) -> int:
    return (~data_out) & 0xFF if flag else data_out & 0xFF


# ---------- bit <-> tensor helpers (LSB-first, +/-1 encoding) ----------
def byte_bits(v: int, n: int = 8) -> torch.Tensor:
    return torch.tensor([1.0 if (v >> k) & 1 else -1.0 for k in range(n)], dtype=torch.float32)


def bits_to_int(bits) -> int:
    return sum((1 << k) for k, b in enumerate(bits) if b > 0)


# ---------- the neural units ----------
class NeuralDBIEncode(nn.Module):
    """8 data bits -> 9 bits (8 encoded data + 1 DBI flag)."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(8, 128), nn.GELU(),
            nn.Linear(128, 128), nn.GELU(),
            nn.Linear(128, 9),
        )

    def forward(self, x):
        return self.net(x)

    @torch.no_grad()
    def encode(self, byte: int) -> tuple[int, int]:
        y = (self(byte_bits(byte).unsqueeze(0))[0] > 0).int().tolist()
        return bits_to_int(y[:8]), int(y[8])


class NeuralDBIDecode(nn.Module):
    """9 bits (8 data + flag) -> 8 recovered data bits."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(9, 128), nn.GELU(),
            nn.Linear(128, 128), nn.GELU(),
            nn.Linear(128, 8),
        )

    def forward(self, x):
        return self.net(x)

    @torch.no_grad()
    def decode(self, data_out: int, flag: int) -> int:
        inp = torch.cat([byte_bits(data_out), torch.tensor([1.0 if flag else -1.0])]).unsqueeze(0)
        y = (self(inp)[0] > 0).int().tolist()
        return bits_to_int(y)


# ---------- exhaustive datasets + verification ----------
def encode_domain():
    X = torch.stack([byte_bits(b) for b in range(256)])
    Y = torch.stack([
        torch.tensor([1.0 if (enc >> k) & 1 else -1.0 for k in range(8)] + [1.0 if fl else -1.0])
        for enc, fl in (golden_encode(b) for b in range(256))
    ])
    return X, Y


def decode_domain():
    X, Y = [], []
    for data_out in range(256):
        for flag in (0, 1):
            X.append(torch.cat([byte_bits(data_out), torch.tensor([1.0 if flag else -1.0])]))
            b = golden_decode(data_out, flag)
            Y.append(torch.tensor([1.0 if (b >> k) & 1 else -1.0 for k in range(8)]))
    return torch.stack(X), torch.stack(Y)


@torch.no_grad()
def verify_encode(unit: NeuralDBIEncode) -> tuple[int, int]:
    ok = 0
    for b in range(256):
        g_enc, g_fl = golden_encode(b)
        n_enc, n_fl = unit.encode(b)
        ok += (n_enc == g_enc and n_fl == g_fl)
    return ok, 256


@torch.no_grad()
def verify_decode(unit: NeuralDBIDecode) -> tuple[int, int]:
    ok = 0
    for data_out in range(256):
        for flag in (0, 1):
            ok += (unit.decode(data_out, flag) == golden_decode(data_out, flag))
    return ok, 512