"""Load the native Qwen3 codec generation graph with explicit tensor accounting."""
from __future__ import annotations
from pathlib import Path
import torch
from safetensors import safe_open
from torch import Tensor, nn
from torch.nn import functional
from transformers import Qwen3Config
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3RotaryEmbedding
from .configuration import Geometry
from .layout import CODEBOOKS, VOCABULARY

class Generator(nn.Module):

    def __init__(self, geometry: Geometry) -> None:
        super().__init__()
        self.geometry = geometry
        config = Qwen3Config(hidden_size=geometry.hidden_size, intermediate_size=geometry.intermediate_size, num_hidden_layers=geometry.layers, num_attention_heads=geometry.heads, num_key_value_heads=geometry.kv_heads, head_dim=geometry.head_dim, vocab_size=geometry.text_vocabulary, rms_norm_eps=geometry.rms_epsilon, rope_parameters={'rope_type': 'default', 'rope_theta': geometry.rope_theta}, max_position_embeddings=32768, attention_dropout=0, use_cache=False)
        self.body = Qwen3Model(config)
        self.audio_embedding = nn.Embedding(CODEBOOKS * VOCABULARY, geometry.hidden_size)

    def embed_audio(self, codes: Tensor) -> Tensor:
        offsets = torch.arange(CODEBOOKS, device=codes.device) * VOCABULARY
        embedded: Tensor = self.audio_embedding(codes.long() + offsets)
        return embedded.sum(dim=-2)

    def enable_activation_checkpointing(self) -> None:
        self.body.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})

    def hidden(self, text_ids: Tensor, codes: Tensor, audio_mask: Tensor, valid: Tensor) -> Tensor:
        embedded: Tensor = self.body.embed_tokens(text_ids)
        embedded = torch.where(audio_mask.unsqueeze(-1), self.embed_audio(codes), embedded)
        output = self.body(inputs_embeds=embedded, attention_mask=valid, use_cache=False)
        if not isinstance(output, BaseModelOutputWithPast):
            message = f'unexpected Qwen3 output: {type(output).__name__}'
            raise TypeError(message)
        hidden = output.last_hidden_state
        if not isinstance(hidden, Tensor):
            message = 'Qwen3 returned no hidden-state tensor'
            raise TypeError(message)
        return hidden

    def classify(self, hidden: Tensor) -> Tensor:
        logits = functional.linear(hidden, self.audio_embedding.weight)
        return logits.reshape(*hidden.shape[:-1], CODEBOOKS, VOCABULARY)

def map_weight(name: str) -> str | None:
    if name.startswith('tied.embedding.modality_embeddings.0.model.'):
        return None
    if name.startswith('tied.head.'):
        return None
    prefixes = (('tied.embedding.text_embedding.', 'body.embed_tokens.'), ('tied.embedding.modality_embeddings.0.embedding.', 'audio_embedding.'), ('body.', 'body.'))
    for source, destination in prefixes:
        if name.startswith(source):
            return destination + name.removeprefix(source)
    message = f'unrecognized native checkpoint tensor {name!r}'
    raise ValueError(message)

def load_generator(path: Path, geometry: Geometry, device: torch.device) -> tuple[Generator, dict[str, object]]:
    with torch.device('meta'):
        model = Generator(geometry)
    expected = model.state_dict()
    seen: set[str] = set()
    dropped: list[str] = []
    with safe_open(path, framework='pt', device='cpu') as checkpoint:
        source_keys = set(checkpoint.keys())
        for source in source_keys:
            destination = map_weight(source)
            if destination is None:
                dropped.append(source)
                continue
            if destination not in expected or destination in seen:
                message = f'unexpected/duplicate tensor {source!r} -> {destination!r} at {path}'
                raise ValueError(message)
            tensor = checkpoint.get_tensor(source)
            if tensor.shape != expected[destination].shape:
                message = f'shape mismatch {path}:{source}: {tensor.shape}'
                message += f' != {expected[destination].shape}'
                raise ValueError(message)
            model.load_state_dict({destination: tensor.to(device)}, strict=False, assign=True)
            seen.add(destination)
        missing = set(expected) - seen
        if missing:
            message = f'missing native tensors at {path}: {sorted(missing)}'
            raise ValueError(message)
        tied_pairs = (('tied.head.text_head.weight', 'tied.embedding.text_embedding.weight'), ('tied.head.modality_heads.0.weight', 'tied.embedding.modality_embeddings.0.embedding.weight'))
        for left, right in tied_pairs:
            if left not in source_keys or right not in source_keys or (not torch.equal(checkpoint.get_tensor(left), checkpoint.get_tensor(right))):
                message = f'missing or unequal tied tensors at {path}: {left}, {right}'
                raise ValueError(message)
    model.body.rotary_emb = Qwen3RotaryEmbedding(model.body.config).to(device)
    model.requires_grad_(False)
    return (model, {'loaded': sorted(seen), 'dropped': dropped, 'ties_verified': True})