"""
ViuMini-Dense-360M — Sovereign 32-Layer Dense Transformer Architecture
Canonical implementation for 366.6M parameter Dense SLM.
- 32 layers: 100% Dense SwiGLU FFN (No MoE routing overhead)
- GQA (15 Q heads, 5 KV heads, head_dim 64, 3:1 ratio)
- RMSNorm + QK-Norm + RoPE (theta 500k) + Hybrid SWA/Full Attention (ratio 3)
- Tied Embeddings (vocab 48,000, dim 960) -> 366,617,536 params
- Native PyTorch SDPA (FlashAttention / memory-efficient kernel)
Run:
  python viu1_dense.py --config ../configs/viu1_dense_config.yaml
"""
import argparse
import math
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    import yaml
except ImportError:
    yaml = None


@dataclass
class ViuDenseArgs:
    dim: int = 960
    n_layers: int = 32
    n_heads: int = 15
    n_kv_heads: int = 5
    vocab_size: int = 48000
    multiple_of: int = 64
    ffn_dim_multiplier: float = 2.7
    norm_eps: float = 1e-5
    rope_theta: float = 500000.0
    max_seq_len: int = 2048
    sliding_window: int = 512
    hybrid_ratio: int = 3  # Every 3rd block full attention, others sliding window
    dropout: float = 0.0
    bias: bool = False
    tie_embeddings: bool = True
    qk_norm: bool = True
    multi_token: bool = False
    # Backward compatibility / config fallback fields
    moe_every: int = 999
    moe_num_experts: int = 0
    moe_top_k: int = 0
    moe_expert_div: int = 1
    moe_shared: bool = False
    moe_balance_weight: float = 0.0
    moe_capacity_factor: float = 1.0
    switch_gate: bool = False
    moe_num_routed_experts: int = None
    moe_shared_experts: int = None


# Backward compatibility aliases
ViuArgs = ViuDenseArgs


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        v = x.pow(2).mean(-1, keepdim=True)
        return self.weight * (x * torch.rsqrt(v + self.eps))


def rope_cos_sin(head_dim: int, seq_len: int, theta: float, device: torch.device, dtype: torch.dtype = torch.float32):
    assert head_dim % 2 == 0
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    t = torch.arange(seq_len, device=device, dtype=torch.float32)
    f = torch.outer(t, inv)
    return torch.cos(f).to(dtype), torch.sin(f).to(dtype)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    d = x.shape[-1]
    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
    cos = cos[: x.shape[1]].unsqueeze(0).unsqueeze(2).to(dtype=x.dtype)
    sin = sin[: x.shape[1]].unsqueeze(0).unsqueeze(2).to(dtype=x.dtype)
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class Attention(nn.Module):
    def __init__(self, a: ViuDenseArgs, sliding: bool):
        super().__init__()
        self.n_heads = a.n_heads
        self.n_kv = a.n_kv_heads
        self.head_dim = a.dim // a.n_heads
        assert a.dim % a.n_heads == 0, f"dim {a.dim} must be divisible by n_heads {a.n_heads}"
        self.n_rep = a.n_heads // a.n_kv_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.sliding = sliding
        self.window = a.sliding_window
        self.qk_norm = a.qk_norm
        self._mask_cache = None

        self.wq = nn.Linear(a.dim, a.n_heads * self.head_dim, bias=a.bias)
        self.wk = nn.Linear(a.dim, self.n_kv * self.head_dim, bias=a.bias)
        self.wv = nn.Linear(a.dim, self.n_kv * self.head_dim, bias=a.bias)
        self.wo = nn.Linear(a.n_heads * self.head_dim, a.dim, bias=a.bias)
        if a.qk_norm:
            self.q_norm = RMSNorm(self.head_dim, a.norm_eps)
            self.k_norm = RMSNorm(self.head_dim, a.norm_eps)
        self.drop = nn.Dropout(a.dropout)

    def _get_mask(self, T: int, device: torch.device):
        if self._mask_cache is not None and self._mask_cache.shape == (T, T) and self._mask_cache.device == device:
            return self._mask_cache
        causal = torch.full((T, T), float("-inf"), device=device)
        causal = torch.triu(causal, diagonal=1)
        if self.sliding:
            win = torch.full((T, T), float("-inf"), device=device)
            idx = torch.arange(T, device=device)
            ok = (idx.unsqueeze(1) - idx.unsqueeze(0) < self.window)
            win = torch.where(ok, torch.zeros((), device=device), win)
            causal = torch.minimum(causal, win)
        self._mask_cache = causal
        return causal

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv, self.head_dim)
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=2)
            v = v.repeat_interleave(self.n_rep, dim=2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        drop_p = self.drop.p if self.training else 0.0
        if self.sliding:
            causal = self._get_mask(T, x.device).to(dtype=q.dtype)
            attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=causal, dropout_p=drop_p)
        else:
            attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop_p)
        attn_out = attn_out.transpose(1, 2).contiguous()
        out = attn_out.view(B, T, -1)
        return self.wo(out)


def ffn_hidden(dim: int, mult: float, multiple_of: int) -> int:
    h = int(2 * 4 * dim / 3)
    if mult:
        h = int(mult * dim)
    return multiple_of * ((h + multiple_of - 1) // multiple_of)


class DenseFFN(nn.Module):
    def __init__(self, a: ViuDenseArgs):
        super().__init__()
        h = ffn_hidden(a.dim, a.ffn_dim_multiplier, a.multiple_of)
        self.w1 = nn.Linear(a.dim, h, bias=a.bias)
        self.w2 = nn.Linear(h, a.dim, bias=a.bias)
        self.w3 = nn.Linear(a.dim, h, bias=a.bias)
        self.drop = nn.Dropout(a.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class Block(nn.Module):
    def __init__(self, a: ViuDenseArgs, idx: int):
        super().__init__()
        sliding = ((idx + 1) % a.hybrid_ratio != 0)
        self.sliding = sliding
        self.attn_norm = RMSNorm(a.dim, a.norm_eps)
        self.attn = Attention(a, sliding=sliding)
        self.ffn_norm = RMSNorm(a.dim, a.norm_eps)
        self.ffn = DenseFFN(a)
        self.is_moe = False  # 100% Dense architecture

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        x = x + self.ffn(self.ffn_norm(x))
        return x, None, None


class Viu1Dense(nn.Module):
    def __init__(self, a: ViuDenseArgs):
        super().__init__()
        self.args = a
        self.tok_emb = nn.Embedding(a.vocab_size, a.dim)
        self.drop = nn.Dropout(a.dropout)
        self.layers = nn.ModuleList([Block(a, i) for i in range(a.n_layers)])
        self.norm = RMSNorm(a.dim, a.norm_eps)
        self.head = nn.Linear(a.dim, a.vocab_size, bias=False)
        self.mt_head = nn.Linear(a.dim, a.vocab_size, bias=False) if a.multi_token else None
        self.use_checkpointing = False
        self.apply(self._init)
        # Tie AFTER init so apply does not overwrite tied weights twice
        if a.tie_embeddings:
            self.head.weight = self.tok_emb.weight

    def gradient_checkpointing_enable(self):
        self.use_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.use_checkpointing = False

    @staticmethod
    def _init(m):
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight, gain=0.7)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, ids: torch.Tensor, targets: torch.Tensor = None):
        B, T = ids.shape
        assert T <= self.args.max_seq_len, f"Sequence length {T} exceeds max_seq_len {self.args.max_seq_len}"
        x = self.drop(self.tok_emb(ids))
        cos, sin = rope_cos_sin(
            self.args.dim // self.args.n_heads, T,
            self.args.rope_theta, ids.device,
        )
        for blk in self.layers:
            if self.use_checkpointing and self.training:
                def _create_custom_forward(module):
                    def _custom_forward(*inputs):
                        return module(*inputs)
                    return _custom_forward
                nx, _, _ = torch.utils.checkpoint.checkpoint(
                    _create_custom_forward(blk), x, cos, sin, use_reentrant=False
                )
            else:
                nx, _, _ = blk(x, cos, sin)
            x = nx

        h = self.norm(x)
        loss, mt_loss, logits = None, None, None
        if targets is not None:
            flat_h = h.reshape(-1, self.args.dim)
            flat_tgt = targets.reshape(-1)
            chunk_size = 1024 if (h.is_cuda and torch.cuda.get_device_properties(h.device).total_memory >= 14 * 1024**3) else 512
            n_tokens = flat_tgt.numel()
            total_loss = 0.0
            for ci in range(0, n_tokens, chunk_size):
                ch_h = flat_h[ci : ci + chunk_size]
                ch_tgt = flat_tgt[ci : ci + chunk_size]
                ch_logits = self.head(ch_h)
                total_loss = total_loss + F.cross_entropy(ch_logits, ch_tgt, reduction="sum")
            loss = total_loss / max(n_tokens, 1)

            if self.mt_head is not None and T > 1:
                flat_h_mt = h[:, :-1, :].reshape(-1, self.args.dim)
                flat_tgt_mt = targets[:, 1:].reshape(-1)
                n_tokens_mt = flat_tgt_mt.numel()
                total_mt_loss = 0.0
                for ci in range(0, n_tokens_mt, chunk_size):
                    ch_h_mt = flat_h_mt[ci : ci + chunk_size]
                    ch_tgt_mt = flat_tgt_mt[ci : ci + chunk_size]
                    ch_logits_mt = self.mt_head(ch_h_mt)
                    total_mt_loss = total_mt_loss + F.cross_entropy(ch_logits_mt, ch_tgt_mt, reduction="sum")
                mt_loss = total_mt_loss / max(n_tokens_mt, 1)
                loss = loss + 0.1 * mt_loss
        else:
            logits = self.head(h)

        aux_info = {
            "moe_aux": 0.0,
            "mt_loss": float(mt_loss.detach()) if isinstance(mt_loss, torch.Tensor) else 0.0,
            "usages": [],
        }
        return logits, loss, aux_info

    def param_counts(self):
        total = sum(p.numel() for p in self.parameters())
        # Pure Dense: 100% of parameters are active in every forward pass
        active = total
        return total, active


# Canonical aliases
Viu1MoE = Viu1Dense


def load_args(path: str) -> ViuDenseArgs:
    if path and yaml and Path(path).exists():
        d = yaml.safe_load(open(path, encoding="utf-8")) or {}
        return ViuDenseArgs(**{k: v for k, v in d.items() if k in ViuDenseArgs.__dataclass_fields__})
    return ViuDenseArgs()


__all__ = [
    "ViuDenseArgs",
    "Viu1Dense",
    "ViuArgs",
    "Viu1MoE",
    "RMSNorm",
    "Attention",
    "DenseFFN",
    "Block",
    "rope_cos_sin",
    "apply_rope",
    "ffn_hidden",
    "load_args",
]


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="../configs/viu1_dense_config.yaml")
    args = ap.parse_args()
    cfg = load_args(args.config)
    m = Viu1Dense(cfg)
    total, active = m.param_counts()
    print(f"[config] dim={cfg.dim} L={cfg.n_layers} H={cfg.n_heads} KV={cfg.n_kv_heads} vocab={cfg.vocab_size}")
    print(f"[params] total={total:,} ({total/1e6:.1f}M) | active={active:,} ({active/1e6:.1f}M)")
    print(f"[target] ViuMini-Dense-360M (32L Dense, 366.6M params, 100% active)")
    m.eval()
    with torch.no_grad():
        ids = torch.randint(0, cfg.vocab_size, (1, 16))
        logits, loss, aux = m(ids, ids)
        logits_str = tuple(logits.shape) if logits is not None else "(loss mode, logits purged)"
        print(f"[smoke] logits {logits_str} loss {loss.item():.3f} aux {aux} OK")
        n_dense = len(m.layers)
        n_sliding = sum(1 for b in m.layers if b.sliding)
        n_full = sum(1 for b in m.layers if not b.sliding)
        print(f"[layers] dense={n_dense} sliding={n_sliding} full={n_full}")