"""Independent CPU E2M1/E8M0 decoder and admission; no CUDA import or access."""
from __future__ import annotations

import torch

E2M1: tuple[float, ...] = (0., .5, 1., 1.5, 2., 3., 4., 6.,
                          -0., -.5, -1., -1.5, -2., -3., -4., -6.)


def decode_codes(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """Float64 oracle, including signed zero, scale-0 and NaN scale-255."""
    if codes.device.type != 'cpu' or scales.device.type != 'cpu':
        raise ValueError('reference decoder requires CPU tensors')
    values = torch.tensor(E2M1, dtype=torch.float64)[codes.long()]
    values = torch.ldexp(values, scales.to(torch.int32) - 127)
    return torch.where(scales == 255, float('nan'), values)


def unpack(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    if packed.dtype != torch.uint8 or scales.dtype != torch.uint8:
        raise ValueError('checkpoint codes and E8M0 scales must be uint8')
    if packed.ndim != 3 or scales.shape != (*packed.shape[:2], packed.shape[2] // 16):
        raise ValueError('expected [E,N,K/2] codes and [E,N,K/32] scales')
    if packed.shape[-1] % 16:
        raise ValueError('K must be a multiple of group size 32')
    codes = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2)
    return decode_codes(codes, scales.repeat_interleave(32, dim=-1))


def admit(packed: torch.Tensor, scales: torch.Tensor) -> dict[str, int]:
    """Bounded load-time validation. Never narrow/repair unsupported native values.

    Inspect one expert at a time so preparing a full bank never expands the
    checkpoint in RAM. BF16 subnormals decode exactly but MMA may flush them:
    reject them for computation rather than asserting unsupported fidelity.
    """
    if packed.device.type != 'cpu' or scales.device.type != 'cpu':
        raise ValueError('admission must run on the CPU before GPU allocation')
    if packed.ndim != 3 or packed.shape[0] == 0 or packed.shape[1] == 0:
        raise ValueError('expected nonempty expert bank')
    result = {'values': 0, 'negative_zeros': 0, 'scale_min': 255, 'scale_max': 0}
    for e in range(packed.shape[0]):
        decoded = unpack(packed[e:e+1], scales[e:e+1])
        if not bool(torch.isfinite(decoded).all()):
            raise ValueError(f'expert {e}: NaN/nonfinite native scale/value')
        widened = decoded.to(torch.bfloat16).to(torch.float64)
        if not torch.equal(decoded, widened):
            raise ValueError(f'expert {e}: native values not exactly representable in BF16')
        subnormal = (decoded != 0) & (decoded.abs() < torch.finfo(torch.bfloat16).tiny)
        if bool(subnormal.any()):
            raise ValueError(f'expert {e}: BF16 subnormal operands require a wider arithmetic path')
        result['values'] += decoded.numel()
        result['negative_zeros'] += int(((decoded == 0) & decoded.signbit()).sum())
        result['scale_min'] = min(result['scale_min'], int(scales[e].min()))
        result['scale_max'] = max(result['scale_max'], int(scales[e].max()))
    return result