"""Hub-layout adapters around Helion's synchronized PyTorch references."""

from __future__ import annotations

import torch

from ._helion_reference import (
    chunked_linear_attn_reference,
    naive_recurrent_reference,
    rel_error,
    recurrent_step_reference,
)


def relative_error(a: torch.Tensor | None, b: torch.Tensor | None) -> float:
    return rel_error(a, b)


def make_inputs(
    device: torch.device,
    *,
    b: int = 1,
    t: int = 64,
    h: int = 2,
    d: int = 32,
    dv: int = 32,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    tensors = [
        torch.randn(b, t, h, dim, device=device, dtype=torch.float32)
        for dim in (d, d, dv)
    ]
    return tuple(x.to(torch.bfloat16) for x in tensors)


def _head_first(x: torch.Tensor) -> torch.Tensor:
    return x.transpose(1, 2).contiguous()


def recurrent_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    g: torch.Tensor | None = None,
    beta: torch.Tensor | None = None,
    scale: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run Helion's recurrent reference on time-first Hub inputs."""
    qh, kh, vh = (_head_first(x) for x in (q, k, v))
    gh = (
        _head_first(g)
        if g is not None
        else torch.zeros(
            q.size(0),
            q.size(2),
            q.size(1),
            device=q.device,
            dtype=torch.float32,
        )
    )
    bh = _head_first(beta) if beta is not None else None

    output = naive_recurrent_reference(
        qh,
        kh,
        vh,
        gh.float(),
        beta=bh,
        q_scale=scale,
    )

    state = torch.zeros(
        q.size(0),
        q.size(2),
        q.size(3),
        v.size(3),
        device=q.device,
        dtype=torch.float32,
    )
    for index in range(q.size(1)):
        decay = gh[:, :, index : index + 1].float().exp()
        beta_value = bh[:, :, index : index + 1].float() if bh is not None else None
        _, state = recurrent_step_reference(
            qh[:, :, index : index + 1].float() * scale,
            kh[:, :, index : index + 1].float(),
            vh[:, :, index : index + 1].float(),
            state,
            alpha=decay,
            beta_val=beta_value,
        )

    return _head_first(output), state


def chunked_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    g: torch.Tensor | None = None,
    beta: torch.Tensor | None = None,
    scale: float,
    chunk_size: int = 64,
) -> torch.Tensor:
    """Run Helion's differentiable chunked reference on Hub-layout inputs."""
    qh, kh, vh = (_head_first(x) for x in (q, k, v))
    gh = (
        _head_first(g)
        if g is not None
        else torch.zeros(
            q.size(0),
            q.size(2),
            q.size(1),
            device=q.device,
            dtype=torch.float32,
        )
    )
    bh = _head_first(beta) if beta is not None else None
    output = chunked_linear_attn_reference(
        qh * scale,
        kh,
        vh,
        gh,
        beta=bh,
        C=chunk_size,
    )
    return _head_first(output)


def assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
    torch.testing.assert_close(actual.float(), expected.float(), atol=6e-2, rtol=3e-2)