from __future__ import annotations

from functools import lru_cache
from functools import partial

import torch
import torch.nn as nn

from .oat_linear import LinearHead, SampleEmbedder
from .oat_transformer import Transformer
from .oat_transformer import Fp32LayerNorm
from .positional import PositionalEmbeddingAdder


@lru_cache(maxsize=32)
def _register_attention_mask(
    action_len: int, register_len: int, device_str: str
) -> torch.Tensor:
    """
    OAT-style mask for SDPA attention (`True` means attended/allowed):
    - action tokens cannot attend to registers
    - register tokens attend to all actions + causal registers
    """
    total = action_len + register_len
    mask = torch.zeros((total, total), dtype=torch.bool, device=device_str)

    # Action queries cannot read any register keys.
    mask[:action_len, action_len:] = True

    # Register queries cannot read future registers.
    reg_future = torch.triu(
        torch.ones((register_len, register_len), dtype=torch.bool, device=device_str),
        diagonal=1,
    )
    mask[action_len:, action_len:] = reg_future
    return ~mask


class RegisterEncoder(nn.Module):
    def __init__(
        self,
        sample_dim: int,
        sample_horizon: int,
        emb_dim: int,
        head_dim: int,
        depth: int,
        pdropout: float,
        latent_dim: int,
        num_registers: int,
    ):
        super().__init__()

        if emb_dim % head_dim != 0:
            raise ValueError(
                f"emb_dim ({emb_dim}) must be divisible by head_dim ({head_dim})"
            )

        self.sample_proj = SampleEmbedder(
            sample_dim, emb_dim, weight_init_style="xavier"
        )
        self.pos_emb = PositionalEmbeddingAdder(emb_dim, sample_horizon)
        self.registers = nn.Parameter(torch.randn(num_registers, emb_dim))

        self.transformer = Transformer(
            dim=emb_dim,
            depth=depth,
            head_dim=head_dim,
            drop=pdropout,
        )
        self.head = LinearHead(
            emb_dim,
            latent_dim,
            weight_init_style="xavier",  # Changed from "zero" to allow gradient flow
            norm_layer=partial(Fp32LayerNorm, bias=False, elementwise_affine=False),
        )
        self.num_registers = num_registers

    def forward(self, actions: torch.Tensor) -> torch.Tensor:
        b, t, _ = actions.shape
        x = self.sample_proj(actions)
        x = self.pos_emb(x)

        regs = self.registers.unsqueeze(0).expand(b, -1, -1)
        x = torch.cat([x, regs], dim=1)

        mask = _register_attention_mask(t, self.num_registers, str(actions.device))
        x = self.transformer(x, block_mask=mask)
        reg_tokens = x[:, t:]
        return self.head(reg_tokens)
