"""
VIU-Transformer — Frontier 48-Layer Mixture of Experts (MoE) Architecture
Total ~1.856B (with MTP) / ~1.819B (base), Active ~300.5M (per token)
- 48 layers: Hierarchical Cognitive Routing (HCR) across 4 cognitive tiers
- 1,392 total micro-experts: 48 layers x (1 permanent shared expert + 28 routed experts, Top-2)
- Multi-Head Latent Attention (MLA): DeepSeek-V3/V4 style low-rank KV compression + decoupled RoPE
- Dual QK-Norm & Hybrid Attention: Sliding Window (512w) + Full Attention every 3rd layer
- Logit Soft-Capping: Attention logits capped at 50.0 (tanh), Output logits capped at 30.0 (tanh)
- RoPE Extended Base Theta: 500,000 (YaRN ready)
- Viu-Lookahead: Multi-Token Prediction (MTP) head for forward planning
Run:
  python viu_moe.py --config ../configs/model_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 ViuArgs:
    dim: int = 768
    n_layers: int = 48
    n_heads: int = 12
    n_kv_heads: int = 2  # GQA fallback / reference alias
    head_dim: int = 64
    rope_dim: int = 64
    q_lora_rank: int = 256
    kv_lora_rank: int = 256
    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 Nth block is full causal attention, others sliding
    dropout: float = 0.0
    bias: bool = False
    tie_embeddings: bool = True
    qk_norm: bool = True  # Gemma 4 / Gemma 3 primary attention stabilizer
    attn_logit_softcapping: float = 0.0  # Gemma 4: pure QK-Norm replaces Gemma 2 tanh softcap
    final_logit_softcapping: float = 0.0  # Gemma 4: pure RMSNorm head
    moe_every: int = 1
    moe_num_experts: int = 28
    moe_num_routed_experts: int = 28
    moe_shared_experts: int = 1
    moe_top_k: int = 2
    moe_expert_div: int = 4  # expert hidden = dense_hidden / div
    moe_shared: bool = True
    moe_balance_weight: float = 0.05
    moe_capacity_factor: float = 1.25  # Monitoring-only capacity factor
    switch_gate: bool = False
    multi_token: bool = True
    use_mhc: bool = True  # DeepSeek-V4.1 Manifold-Constrained Hyper-Connections
    router_bias: bool = True  # DeepSeek-V4.1 Dynamic Auxiliary-Loss-Free Router Bias

    def __post_init__(self):
        if self.dim % self.n_heads == 0:
            self.head_dim = self.dim // self.n_heads
        if self.rope_dim > self.head_dim:
            self.rope_dim = self.head_dim
        if self.q_lora_rank > self.dim:
            self.q_lora_rank = min(256, self.dim)
        if self.kv_lora_rank > self.dim:
            self.kv_lora_rank = min(256, self.dim)
        if self.moe_num_routed_experts != self.moe_num_experts:
            if self.moe_num_experts != 28 and self.moe_num_routed_experts == 28:
                self.moe_num_routed_experts = self.moe_num_experts
            elif self.moe_num_routed_experts != 28 and self.moe_num_experts == 28:
                self.moe_num_experts = self.moe_num_routed_experts
            else:
                self.moe_num_routed_experts = self.moe_num_experts
        if not self.moe_shared:
            self.moe_shared_experts = 0


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

    def forward(self, x):
        orig_dtype = x.dtype
        x_f = x.float()
        v = x_f.pow(2).mean(-1, keepdim=True)
        return (x_f * torch.rsqrt(v + self.eps)).to(orig_dtype) * self.weight


def rope_cos_sin(rope_dim, seq_len, theta, device, dtype=torch.float32):
    assert rope_dim % 2 == 0
    inv = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, device=device, dtype=torch.float32) / rope_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, cos, sin):
    # x shape: [B, T, H, rope_dim] or [B, T, 1, rope_dim]
    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 MultiHeadLatentAttention(nn.Module):
    """
    Multi-Head Latent Attention (MLA) — DeepSeek-V3/V4 style + Gemma-2 Softcapping:
    - Low-rank query compression (q_lora_rank) + decoupled RoPE
    - Low-rank key/value compression (kv_lora_rank) + decoupled RoPE
    - Dual QK-Norm
    - Logit soft-capping (tanh capper at 50.0)
    - Hybrid Multi-Scale Sliding Window (512w / Full)
    """
    def __init__(self, a: ViuArgs, sliding: bool):
        super().__init__()
        self.dim = a.dim
        self.n_heads = a.n_heads
        self.head_dim = a.head_dim
        self.rope_dim = a.rope_dim
        self.q_lora_rank = a.q_lora_rank
        self.kv_lora_rank = a.kv_lora_rank
        self.sliding = sliding
        self.window = a.sliding_window
        self.softcap = a.attn_logit_softcapping
        self.qk_norm = a.qk_norm
        self.scale = 1.0 / math.sqrt(self.head_dim + self.rope_dim)
        self._mask_cache = None

        # Query compression
        self.wq_down = nn.Linear(a.dim, a.q_lora_rank, bias=a.bias)
        self.q_norm = RMSNorm(a.q_lora_rank, a.norm_eps)
        self.wq_up = nn.Linear(a.q_lora_rank, a.n_heads * a.head_dim, bias=a.bias)
        self.wq_rope = nn.Linear(a.q_lora_rank, a.n_heads * a.rope_dim, bias=a.bias)

        # KV compression
        self.wkv_down = nn.Linear(a.dim, a.kv_lora_rank, bias=a.bias)
        self.kv_norm = RMSNorm(a.kv_lora_rank, a.norm_eps)
        self.wkv_up_k = nn.Linear(a.kv_lora_rank, a.n_heads * a.head_dim, bias=a.bias)
        self.wkv_up_v = nn.Linear(a.kv_lora_rank, a.n_heads * a.head_dim, bias=a.bias)
        self.wk_rope = nn.Linear(a.dim, a.rope_dim, bias=a.bias)

        # Output projection
        self.wo = nn.Linear(a.n_heads * a.head_dim, a.dim, bias=a.bias)

        if a.qk_norm:
            self.q_head_norm = RMSNorm(a.head_dim, a.norm_eps)
            self.k_head_norm = RMSNorm(a.head_dim, a.norm_eps)
        else:
            self.q_head_norm = None
            self.k_head_norm = None

        self.drop = nn.Dropout(a.dropout)

    def _get_mask(self, T, 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:
            idx = torch.arange(T, device=device)
            ok = (idx.unsqueeze(1) - idx.unsqueeze(0) < self.window)
            causal = torch.where(ok, causal, torch.tensor(float("-inf"), device=device))
        self._mask_cache = causal
        return causal

    def forward(self, x, cos, sin):
        B, T, _ = x.shape

        # 1. Query path
        cq = self.q_norm(self.wq_down(x))
        q = self.wq_up(cq).view(B, T, self.n_heads, self.head_dim)
        q_rope = self.wq_rope(cq).view(B, T, self.n_heads, self.rope_dim)

        # 2. KV path
        ckv = self.kv_norm(self.wkv_down(x))
        k = self.wkv_up_k(ckv).view(B, T, self.n_heads, self.head_dim)
        v = self.wkv_up_v(ckv).view(B, T, self.n_heads, self.head_dim)
        k_rope = self.wk_rope(x).unsqueeze(2)  # [B, T, 1, rope_dim]
        k_rope = k_rope.expand(B, T, self.n_heads, self.rope_dim)

        if self.q_head_norm is not None:
            q = self.q_head_norm(q)
            k = self.k_head_norm(k)

        # 3. Decoupled RoPE
        q_rope = apply_rope(q_rope, cos, sin)
        k_rope = apply_rope(k_rope, cos, sin)

        # 4. Form full Q and K
        q_full = torch.cat([q, q_rope], dim=-1).transpose(1, 2)  # [B, H, T, head_dim + rope_dim]
        k_full = torch.cat([k, k_rope], dim=-1).transpose(1, 2)  # [B, H, T, head_dim + rope_dim]
        v_full = v.transpose(1, 2)                               # [B, H, T, head_dim]

        # 5. Attention scores & FlashAttention / SDPA Acceleration
        if self.softcap == 0:
            drop_p = self.drop.p if self.training else 0.0
            if not self.sliding:
                # Full causal attention -> hardware-accelerated FlashAttention-2 / SDPA
                out = F.scaled_dot_product_attention(
                    q_full, k_full, v_full,
                    is_causal=True,
                    dropout_p=drop_p,
                    scale=self.scale
                )
            else:
                # Sliding window attention -> SDPA with causal sliding mask
                causal = self._get_mask(T, x.device)
                out = F.scaled_dot_product_attention(
                    q_full, k_full, v_full,
                    attn_mask=causal.unsqueeze(0).unsqueeze(0),
                    dropout_p=drop_p,
                    scale=self.scale
                )
        else:
            # Manual path with soft-capping
            scores = torch.matmul(q_full, k_full.transpose(-1, -2)) * self.scale
            scores = self.softcap * torch.tanh(scores / self.softcap)
            causal = self._get_mask(T, x.device)
            scores = scores + causal.unsqueeze(0).unsqueeze(0)
            probs = F.softmax(scores, dim=-1)
            if self.drop.p > 0 and self.training:
                probs = self.drop(probs)
            out = torch.matmul(probs, v_full)  # [B, H, T, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.wo(out)


def ffn_hidden(dim, mult, multiple_of):
    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: ViuArgs):
        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):
        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class MicroExpert(nn.Module):
    def __init__(self, a: ViuArgs):
        super().__init__()
        h = max(a.multiple_of, ffn_hidden(a.dim, a.ffn_dim_multiplier, a.multiple_of) // a.moe_expert_div)
        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)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class NextGenMoELayer(nn.Module):
    """
    DeepSeek-V4.1 Sparse Mixture of Experts:
    - 1 Permanent Shared Expert (captures universal syntax & common n-grams)
    - 28 Fine-Grained Routed Micro-Experts (Top-2 routing with dynamic router bias)
    - Dynamic router bias for auxiliary-loss-free load balancing (DeepSeek-V4.1)
    - Top-k aware GShard-style auxiliary load-balancing loss
    """
    def __init__(self, a: ViuArgs):
        super().__init__()
        self.k = a.moe_top_k
        self.n = a.moe_num_routed_experts
        self.cap = a.moe_capacity_factor
        self.bal_w = a.moe_balance_weight
        self.router = nn.Linear(a.dim, self.n, bias=False)
        self.router_bias = nn.Parameter(torch.zeros(self.n)) if getattr(a, "router_bias", True) else None
        self.experts = nn.ModuleList([MicroExpert(a) for _ in range(self.n)])
        self.shared_experts = nn.ModuleList([MicroExpert(a) for _ in range(a.moe_shared_experts)]) if a.moe_shared_experts > 0 else None
        self.drop = nn.Dropout(a.dropout)

    def forward(self, x):
        B, T, D = x.shape
        logits = self.router(x)  # [B, T, E]
        if self.router_bias is not None:
            logits = logits + self.router_bias
        probs = F.softmax(logits, dim=-1)
        top_p, top_i = torch.topk(probs, self.k, dim=-1)
        top_p = top_p / (top_p.sum(-1, keepdim=True) + 1e-9)

        # Auxiliary balance loss (top-k aware, GShard-style, computed only if bal_w > 0):
        if self.bal_w > 0:
            with torch.no_grad():
                one_hot = F.one_hot(top_i, self.n).float()  # [B, T, K, E]
                frac = one_hot.sum(-2).mean(dim=1) / max(self.k, 1)  # [B, E]
            mean_p = probs.mean(dim=1)  # [B, E]
            aux_loss = (self.n * (mean_p * frac).sum(-1)).mean() * self.bal_w
        else:
            aux_loss = None

        out = torch.zeros_like(x)
        flat_x = x.reshape(-1, D)
        flat_top_i = top_i.reshape(-1, self.k)
        flat_top_p = top_p.reshape(-1, self.k)
        flat_out = out.reshape(-1, D)
        N = flat_x.shape[0]

        # Ultra-fast vectorized dispatch: 1 argsort + 1 bincount replaces 28x sequential torch.where searches
        token_idx = torch.arange(N, device=x.device).unsqueeze(1).expand(-1, self.k).reshape(-1)
        expert_ids = flat_top_i.reshape(-1)
        weights = flat_top_p.reshape(-1, 1).to(flat_x.dtype)

        sort_order = torch.argsort(expert_ids)
        sorted_tokens = token_idx[sort_order]
        sorted_weights = weights[sort_order]
        counts = torch.bincount(expert_ids, minlength=self.n)

        offsets = 0
        for ei, expert in enumerate(self.experts):
            cnt = counts[ei].item()
            if cnt == 0:
                continue
            sel_tokens = sorted_tokens[offsets : offsets + cnt]
            w = sorted_weights[offsets : offsets + cnt]
            y_sel = expert(flat_x[sel_tokens])
            flat_out.index_add_(0, sel_tokens, (w * y_sel).to(flat_out.dtype))
            offsets += cnt

        out = flat_out.view(B, T, D)
        out = self.drop(out)

        # 1 Permanent Shared Expert unconditionally processes all tokens
        if self.shared_experts is not None:
            for s_exp in self.shared_experts:
                out = out + s_exp(x)

        usage = counts.float() / max(expert_ids.numel(), 1)
        return out, aux_loss, usage


class ManifoldHyperConnection(nn.Module):
    """
    DeepSeek-V4 / V4.1 Manifold-Constrained Hyper-Connection (mHC).
    Replaces standard unconstrained residual x + F(x) with doubly stochastic manifold projection:
      weights = softmax([gate_res, gate_post]) * 2.0
      output = weights[0] * x + weights[1] * F(x)
    Initialized at [0, 0] -> exact identity residual (1.0*x + 1.0*F(x)).
    Prevents representational collapse across 48 ultra-deep layers.
    """
    def __init__(self, dim: int):
        super().__init__()
        self.gate = nn.Parameter(torch.zeros(2))

    def forward(self, x, fx):
        w = (F.softmax(self.gate, dim=0) * 2.0).to(x.dtype)
        return w[0] * x + w[1] * fx


# Alias DeepSeek-V4.1 Compressed Sparse Attention 2 (CSA2)
CompressedSparseAttention2 = MultiHeadLatentAttention


class Block(nn.Module):
    def __init__(self, a: ViuArgs, idx: int):
        super().__init__()
        # hybrid: har Nth block full, baaki sliding
        sliding = ((idx + 1) % a.hybrid_ratio != 0)
        self.sliding = sliding
        self.attn_norm = RMSNorm(a.dim, a.norm_eps)
        self.attn = CompressedSparseAttention2(a, sliding=sliding)
        self.ffn_norm = RMSNorm(a.dim, a.norm_eps)
        self.is_moe = ((idx + 1) % a.moe_every == 0)
        self.ffn = NextGenMoELayer(a) if self.is_moe else DenseFFN(a)
        self.use_mhc = getattr(a, "use_mhc", True)
        if self.use_mhc:
            self.attn_mhc = ManifoldHyperConnection(a.dim)
            self.ffn_mhc = ManifoldHyperConnection(a.dim)
        else:
            self.attn_mhc = None
            self.ffn_mhc = None

    def forward(self, x, cos, sin):
        fx_attn = self.attn(self.attn_norm(x), cos, sin)
        x = self.attn_mhc(x, fx_attn) if self.attn_mhc is not None else x + fx_attn
        if self.is_moe:
            f, aux, usage = self.ffn(self.ffn_norm(x))
            x = self.ffn_mhc(x, f) if self.ffn_mhc is not None else x + f
            return x, aux, usage
        fx_ffn = self.ffn(self.ffn_norm(x))
        x = self.ffn_mhc(x, fx_ffn) if self.ffn_mhc is not None else x + fx_ffn
        return x, None, None


class Viu1MoE(nn.Module):
    def __init__(self, a: ViuArgs):
        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.final_softcap = a.final_logit_softcapping
        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.gradient_checkpointing = False
        self.apply(self._init)
        # tie AFTER init — warna apply() tied weight ko do-baar overwrite karta hai
        if a.tie_embeddings:
            self.head.weight = self.tok_emb.weight

    def gradient_checkpointing_enable(self):
        self.gradient_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.gradient_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, targets=None):
        B, T = ids.shape
        assert T <= self.args.max_seq_len
        x = self.drop(self.tok_emb(ids))
        cos, sin = rope_cos_sin(
            self.args.rope_dim, T,
            self.args.rope_theta, ids.device,
        )
        total_aux, usages = 0.0, []
        use_ckpt = self.gradient_checkpointing and self.training
        for blk in self.layers:
            if use_ckpt:
                def _create_custom_forward(module):
                    def _custom_forward(*inputs):
                        return module(*inputs)
                    return _custom_forward
                nx, aux, usage = torch.utils.checkpoint.checkpoint(
                    _create_custom_forward(blk), x, cos, sin, use_reentrant=False
                )
            else:
                nx, aux, usage = blk(x, cos, sin)
            x = nx
            if aux is not None:
                total_aux = total_aux + aux
                usages.append(usage.detach())
        h = self.norm(x)
        loss, mt_loss = None, None
        if targets is not None:
            # Chunked cross entropy: computes head projection + loss in 256-token chunks
            # Eliminates 300MB+ peak allocation at vocabulary projection on 16GB GPUs
            flat_h = h.reshape(-1, self.args.dim)
            flat_tgt = targets.reshape(-1)
            chunk_size = 2048 if (h.is_cuda and torch.cuda.get_device_properties(h.device).total_memory > 16 * 1024**3) else 256
            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)
                if self.final_softcap > 0:
                    ch_logits = self.final_softcap * torch.tanh(ch_logits / self.final_softcap)
                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)
                    if self.final_softcap > 0:
                        ch_logits_mt = self.final_softcap * torch.tanh(ch_logits_mt / self.final_softcap)
                    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

            if isinstance(total_aux, torch.Tensor):
                loss = loss + total_aux
            logits = None
        else:
            logits = self.head(h)
            if self.final_softcap > 0:
                logits = self.final_softcap * torch.tanh(logits / self.final_softcap)
        aux_info = {
            "moe_aux": total_aux.detach() if isinstance(total_aux, torch.Tensor) else 0.0,
            "mt_loss": mt_loss.detach() if isinstance(mt_loss, torch.Tensor) else 0.0,
            "usages": usages,
        }
        return logits, loss, aux_info

    def param_counts(self):
        total = sum(p.numel() for p in self.parameters())
        dense_like = 0
        for blk in self.layers:
            for p in blk.attn.parameters():
                dense_like += p.numel()
            for p in blk.attn_norm.parameters():
                dense_like += p.numel()
            for p in blk.ffn_norm.parameters():
                dense_like += p.numel()
            if hasattr(blk, "attn_mhc") and blk.attn_mhc is not None:
                for p in blk.attn_mhc.parameters():
                    dense_like += p.numel()
            if hasattr(blk, "ffn_mhc") and blk.ffn_mhc is not None:
                for p in blk.ffn_mhc.parameters():
                    dense_like += p.numel()
            if blk.is_moe:
                moe = blk.ffn
                for p in moe.router.parameters():
                    dense_like += p.numel()
                if hasattr(moe, "router_bias") and moe.router_bias is not None:
                    dense_like += moe.router_bias.numel()
                # k routed experts active per token
                for ei in range(min(moe.k, len(moe.experts))):
                    for p in moe.experts[ei].parameters():
                        dense_like += p.numel()
                # 1 permanent shared expert always active
                if moe.shared_experts is not None:
                    for s_exp in moe.shared_experts:
                        for p in s_exp.parameters():
                            dense_like += p.numel()
            else:
                for p in blk.ffn.parameters():
                    dense_like += p.numel()
        active = dense_like + self.tok_emb.weight.numel() + self.norm.weight.numel()
        if not self.args.tie_embeddings:
            active += self.head.weight.numel()
        return total, active


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


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="../configs/model_config.yaml")
    args = ap.parse_args()
    cfg = load_args(args.config)
    m = Viu1MoE(cfg)
    total, active = m.param_counts()
    print(f"[config] dim={cfg.dim} L={cfg.n_layers} H={cfg.n_heads} "
          f"MLA(q_lora={cfg.q_lora_rank}, kv_lora={cfg.kv_lora_rank}) "
          f"MoE(routed={cfg.moe_num_routed_experts}, shared={cfg.moe_shared_experts}, topk={cfg.moe_top_k})")
    print(f"[params] total={total:,} ({total/1e6:.1f}M / {total/1e9:.3f}B) | active={active:,} ({active/1e6:.1f}M)")
    print(f"[target] total ~1.856B, active ~300.5M (48L NEXT-GEN VIU-TRANSFORMER LOCK)")
    m.eval()
    with torch.no_grad():
        ids = torch.randint(0, cfg.vocab_size, (1, 16))
        logits, loss, aux = m(ids, ids)
        logits_info = "None (freed — targets given)" if logits is None else str(tuple(logits.shape))
        print(f"[smoke] logits={logits_info} loss={loss.item():.3f} aux={aux} OK")
        n_moe = sum(1 for b in m.layers if b.is_moe)
        print(f"[layers] MoE={n_moe} dense={cfg.n_layers - n_moe} "
              f"sliding={sum(1 for b in m.layers if b.sliding)} full={sum(1 for b in m.layers if not b.sliding)}")