"""
Test TinyV4 base model
— load dari HuggingFace Hub (ukung/tinyv4) atau dari folder lokal
— forward pass, generate text (ID & EN)
"""
import torch
import json
import os
import sys

# ---------------------------------------------------------------------------
# 0. Config — ganti ke False kalau mau test dari folder lokal
# ---------------------------------------------------------------------------
USE_HUB = True
HF_REPO = "ukung/tinyv4"

if USE_HUB:
    # Load dari HuggingFace Hub (trust_remote_code=True)
    from transformers import AutoTokenizer, AutoModel
    tokenizer = AutoTokenizer.from_pretrained(HF_REPO)
    model = AutoModel.from_pretrained(HF_REPO, trust_remote_code=True)
    model.head.weight = model.embed.weight  # tie embeddings
    model.eval()
else:
    # Load dari folder lokal
    MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, MODEL_DIR)
    from modeling_tinyv4 import TinyV4, TinyV4Config
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
    model = TinyV4.from_pretrained(MODEL_DIR)
    model.head.weight = model.embed.weight  # tie embeddings
    model.eval()

n_params = sum(p.numel() for p in model.parameters())
print(f"✅ Model loaded: {n_params:,} params ({n_params/1e6:.2f}M)")

# ---------------------------------------------------------------------------
# 2. Config info
# ---------------------------------------------------------------------------
cfg = model.config
print(f"✅ Config: dim={cfg.dim}, depth={cfg.depth}, vocab={cfg.vocab_size}")
print(f"   MoE: {cfg.n_routed} routed + {cfg.n_shared} shared, {cfg.n_active} active")
print(f"   MTP: depth={cfg.mtp_depth}, max_len={cfg.max_len}")

# ---------------------------------------------------------------------------
# 3. Tie check
# ---------------------------------------------------------------------------
assert model.head.weight.data_ptr() == model.embed.weight.data_ptr(), "❌ Embedding tie FAILED!"
print("✅ Embedding tie: OK")

# ---------------------------------------------------------------------------
# 4. Forward pass (smoke test)
# ---------------------------------------------------------------------------
dummy = torch.randint(0, cfg.vocab_size, (2, 64))
with torch.no_grad():
    logits, mtp, bal = model(dummy)

has_nan = torch.isnan(logits).any().item()
has_inf = torch.isinf(logits).any().item()
print(f"✅ Forward pass: logits={logits.shape}, NaN={has_nan}, Inf={has_inf}")
if mtp is not None:
    print(f"   MTP logits: {mtp.shape}, NaN={torch.isnan(mtp).any().item()}")
print(f"   Balance loss: {bal.item():.6f}" if bal is not None else "   Balance loss: None")

# ---------------------------------------------------------------------------
# 5. Generate text
# ---------------------------------------------------------------------------
@torch.no_grad()
def generate(prompt, max_new_tokens=60, temperature=0.8, top_k=40):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    for _ in range(max_new_tokens):
        idx = input_ids[:, -cfg.max_len:]
        logits, _, _ = model(idx)
        logits = logits[:, -1, :] / temperature

        # Top-k filter
        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
        logits[logits < v[:, [-1]]] = float("-inf")

        probs = torch.softmax(logits, dim=-1)

        # Cek NaN di probs — fallback ke uniform
        if torch.isnan(probs).any() or torch.isinf(probs).any():
            probs = torch.ones_like(probs) / probs.size(-1)

        next_token = torch.multinomial(probs, 1)
        input_ids = torch.cat([input_ids, next_token], dim=1)

        if next_token.item() == tokenizer.eos_token_id:
            break

    return tokenizer.decode(input_ids[0], skip_special_tokens=True)

print()
print("=" * 60)
print("📝 GENERATION TEST")
print("=" * 60)

prompts = [
    ("EN", "Once upon a time,"),
    ("EN", "There was a little"),
    ("EN", "In a small village,"),
    ("ID", "Pada suatu hari,"),
    ("ID", "Di sebuah desa kecil,"),
    ("ID", "Alkisah, tersebutlah"),
]

for lang, prompt in prompts:
    output = generate(prompt, max_new_tokens=50, temperature=0.8)
    print(f"  [{lang}] {prompt}")
    print(f"  → {output}")
    print()

print("=" * 60)
print("✅ ALL TESTS PASSED")
print("=" * 60)