"""
Step 2: ADDR_MAP -- physical address -> (rank/bank-group/bank/row/column).

DDR performance depends on which address bits pick the bank vs the row/column,
because interleaving consecutive addresses across banks/bank-groups exposes
parallelism. That mapping is a fixed bit-permutation -> a finite function we can
verify N/N, and it is exactly the kind of per-generation policy the bridge swaps.

Small but representative geometry (12-bit address, 4096 configs):
  column = 4 bits, bank = 2, bank-group = 2, row = 4
laid out as an interleave (bank/bg taken from interior bits for parallelism).
"""
from __future__ import annotations
import torch
from .common import bits_of, int_of, pm

ADDR_BITS = 12
# output bit index -> which input address bit feeds it (the addressing scheme)
PERM = [0, 3, 6, 9,   1, 4,   7, 10,   2, 5, 8, 11]   # [col4][bank2][bg2][row4]
FIELDS = [("col", 4), ("bank", 2), ("bg", 2), ("row", 4)]


def golden_map_bits(addr: int) -> list[int]:
    ib = bits_of(addr, ADDR_BITS)
    return [ib[PERM[o]] for o in range(ADDR_BITS)]


def decode_fields(outbits) -> dict:
    res, i = {}, 0
    for name, w in FIELDS:
        res[name] = int_of(outbits[i:i + w]); i += w
    return res


def domain():
    X = torch.stack([pm(bits_of(a, ADDR_BITS)) for a in range(2 ** ADDR_BITS)])
    Y = torch.tensor([[float(b) for b in golden_map_bits(a)] for a in range(2 ** ADDR_BITS)])
    return X, Y


def run(net, addr: int) -> dict:
    y = (net(pm(bits_of(addr, ADDR_BITS)).unsqueeze(0))[0] > 0).int().tolist()
    return decode_fields(y)