"""
SAD Trainer – Block-AR vectorized training.

Per training step:
  1. Get clean embeddings: clean_embs = leaf_embed[input_ids]  [B, L, d]
  2. Sample per-token levels uniformly in {0..mask_level}
  3. Build noisy_embs from AncestorTable LUT
  4. model.forward(noisy_embs) → leaf_logits [B, L, V]
  5. Loss = L_leaf + lambda_ancestor * L_ancestor
"""

from typing import Dict, Tuple

import torch
import torch.nn as nn

from src.diffusion.noisy_state import NoisyStateBuilder
from src.diffusion.ancestor_table import AncestorTable
from src.losses.sad_loss import SADLoss


class SADTrainer(nn.Module):
    def __init__(
        self,
        config,
        model: nn.Module,
        ancestor_table: AncestorTable,
        loss_fn: SADLoss,
        tokenizer,
        noisy_builder: NoisyStateBuilder,
        dtype=None,
        # kept for API compat (ignored):
        hierarchy=None,
        schedule=None,
    ):
        super().__init__()
        self.config = config
        self.model = model
        self.ancestor_table = ancestor_table
        self.loss_fn = loss_fn
        self.tokenizer = tokenizer
        self.noisy_builder = noisy_builder
        self.dtype = dtype or torch.bfloat16

        device = next(model.parameters()).device
        self.register_buffer("_device_probe", torch.zeros(1, device=device))

    @property
    def device(self):
        return self._device_probe.device

    def forward(
        self,
        batch: Dict[str, torch.Tensor],
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        input_ids = batch["input_ids"].to(self.device)
        attention_mask = batch["attention_mask"].to(self.device)
        B, L = input_ids.shape

        autocast_device = "cuda" if self.device.type == "cuda" else "cpu"

        with torch.autocast(device_type=autocast_device, dtype=self.dtype):
            leaf_emb = self.model.get_leaf_embeddings()              # [V, d]
            mask_emb = leaf_emb[self.tokenizer.mask_token_id]        # [d]

            # num_total_states = num_ancestor_levels + 1 (clean) + 1 (mask)
            num_total_states = self.ancestor_table.num_levels + 2
            levels = self.noisy_builder.sample_levels_uniform(
                B, L, num_total_states, self.device
            )

            noisy_embs, ancestor_log_probs, ancestor_probs_per_lvl, corrupt_mask = \
                self.noisy_builder.build_noisy_embeddings(
                    input_ids, levels, self.ancestor_table, leaf_emb, mask_emb
                )

            fwd_mask = attention_mask if (attention_mask == 0).any() else None
            leaf_logits, _ = self.model.forward(
                noisy_embs,
                attention_mask=fwd_mask,
            )  # [B, L, V]

            loss, metrics = self.loss_fn(
                leaf_logits=leaf_logits,
                input_ids=input_ids,
                levels=levels,
                attention_mask=attention_mask,
                ancestor_log_probs=ancestor_log_probs,
                ancestor_probs_per_level=ancestor_probs_per_lvl,
                corrupt_mask=corrupt_mask,
            )

        return loss, metrics