"""
Nova 2.0 — Mixture-of-Experts (MoE) & Multi-Modal Vision GPU Training Script
Trains 522.8M Nova 2.0 MoE-Vision Model on AMD ROCm GPU (96 GB VRAM)
with Top-2 Sparse Expert Routing and MoE Load Balancing Loss.
"""

import os
import sys
import argparse
import random
import torch
import torch.nn as nn

# Add project root to sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from config.model_config import Nova1Config
from src.tokenizer.bpe_tokenizer import Nova1Tokenizer
from src.model.nova2_moe_vision import Nova2MoEVision
from src.dataset.hf_dataset import create_hf_dataloader


def main():
    parser = argparse.ArgumentParser(description="Pre-train Nova 2.0 MoE-Vision Model")
    parser.add_argument("--tokenizer_path", type=str, default="checkpoints/nova1_gemini_tokenizer.json", help="Tokenizer JSON")
    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size per step")
    parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate")
    parser.add_argument("--num_experts", type=int, default=8, help="Number of MoE Experts")
    parser.add_argument("--top_k", type=int, default=2, help="Top-K active experts per token")
    parser.add_argument("--max_samples", type=int, default=15000, help="Max dataset samples")
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN")
    tokenizer = Nova1Tokenizer.load(args.tokenizer_path)
    print(f"Loaded 32K Tokenizer from '{args.tokenizer_path}'.", flush=True)

    # Nova 2.0 MoE Config (~522.8M Parameters)
    config = Nova1Config(
        vocab_size=tokenizer.vocab_size,
        d_model=768,
        n_heads=12,
        d_ff=3072,
        n_layers_L=4,
        n_layers_H=4,
        max_seq_len=512,
        learning_rate=args.lr,
        device="cuda" if torch.cuda.is_available() else "cpu",
        dtype="bfloat16"
    )

    print(f"\n=======================================================", flush=True)
    print(f"🚀 Building Nova 2.0 Sparse Mixture-of-Experts (MoE) Vision Engine", flush=True)
    print(f"Device: {config.device.upper()} | Experts: {args.num_experts} (Top-{args.top_k}) | Precision: {config.dtype}", flush=True)
    print(f"=======================================================\n", flush=True)

    model = Nova2MoEVision(config, num_experts=args.num_experts, top_k=args.top_k, vision_dim=768).to(config.device)
    num_params = sum(p.numel() for p in model.parameters())
    print(f"🌟 Nova 2.0 MoE Total Parameters: {num_params:,} (~{num_params / 1e6:.1f} Million)", flush=True)

    # Multi-Domain DataLoader
    dataloader = create_hf_dataloader(
        dataset_name="HuggingFaceH4/ultrachat_200k",
        tokenizer=tokenizer,
        max_seq_len=config.max_seq_len,
        batch_size=args.batch_size,
        split="train_sft",
        max_samples=args.max_samples
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_id)

    print(f"\nStarting Nova 2.0 MoE Pre-Training on AMD ROCm GPU across {args.max_samples:,} samples...", flush=True)

    for epoch in range(args.epochs):
        model.train()
        total_loss = 0.0
        total_aux = 0.0
        steps_per_epoch = len(dataloader)

        for step, (b_inp, b_tgt) in enumerate(dataloader):
            b_inp = b_inp.to(config.device)
            b_tgt = b_tgt.to(config.device)

            # Generate synthetic multi-modal vision features (16 patches x 768 dim) for 20% of batches
            pixel_inp = torch.randn(b_inp.shape[0], 16, 768, device=config.device) if random.random() < 0.2 else None

            optimizer.zero_grad()
            with torch.amp.autocast("cuda", enabled=(config.device == "cuda"), dtype=config.get_torch_dtype()):
                logits, _, _, aux_loss = model(b_inp, pixel_features=pixel_inp)
                
                # Align sequence length if vision features prepended patches
                if logits.shape[1] > b_tgt.shape[1]:
                    target_len = b_tgt.shape[1]
                    logits_text = logits[:, -target_len:, :]
                else:
                    logits_text = logits

                loss_lm = criterion(logits_text.reshape(-1, config.vocab_size), b_tgt.reshape(-1))
                # Total Loss = Language Modeling Loss + 0.01 * MoE Load Balancing Loss
                total_step_loss = loss_lm + 0.01 * aux_loss

            total_step_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            total_loss += loss_lm.item()
            total_aux += aux_loss.item()

            if (step + 1) % 25 == 0 or (step + 1) == steps_per_epoch:
                print(f"MoE Epoch {epoch+1}/{args.epochs} [{step+1}/{steps_per_epoch}] — LM Loss: {loss_lm.item():.4f} | MoE Aux Loss: {aux_loss.item():.2f}", flush=True)

        avg_lm_loss = total_loss / steps_per_epoch
        print(f"MoE Epoch {epoch+1}/{args.epochs} Complete — Avg LM Loss: {avg_lm_loss:.4f}\n", flush=True)

        # Save Checkpoint
        ckpt_path = f"checkpoints/nova2_moe_epoch_{epoch+1}.pt"
        torch.save({"model_state": model.state_dict(), "config": config}, ckpt_path)

    final_path = "checkpoints/nova2_moe_final.pt"
    torch.save({"model_state": model.state_dict(), "config": config}, final_path)
    print(f"🚀 Nova 2.0 MoE Pre-Training Complete! Model saved to '{final_path}'.", flush=True)


if __name__ == "__main__":
    main()