"""
Nova 1.0 — Gemini-Style Instruction & Reasoning Pre-Training Script
Trains high-capacity Nova 1.0 model on UltraChat / OpenHermes datasets using AMD ROCm GPU.
"""

import os
import sys
import argparse
import torch

# 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.nova1_hrm import Nova1HRM
from src.dataset.hf_dataset import create_hf_dataloader
from src.trainer.train_loop import Nova1Trainer


def main():
    parser = argparse.ArgumentParser(description="Pre-train Gemini-Style Nova 1.0 Model")
    parser.add_argument("--dataset", type=str, default="HuggingFaceH4/ultrachat_200k", help="HuggingFace dataset name")
    parser.add_argument("--subset", type=str, default=None, help="Dataset subset/config")
    parser.add_argument("--vocab_size", type=int, default=32768, help="Vocabulary size")
    parser.add_argument("--d_model", type=int, default=768, help="Model hidden dimension (191M-350M scale)")
    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=16, help="Batch size per step")
    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
    parser.add_argument("--max_samples", type=int, default=10000, help="Number of dataset samples")
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN")
    tokenizer_path = "checkpoints/nova1_gemini_tokenizer.json"

    # 1. Load 32K Tokenizer
    if not os.path.exists(tokenizer_path):
        print(f"Tokenizer file not found at '{tokenizer_path}'. Please run tokenizer training first.", flush=True)
        return
    tokenizer = Nova1Tokenizer.load(tokenizer_path)
    print(f"Loaded existing 32K Tokenizer from '{tokenizer_path}'.", flush=True)

    # 2. Configure Scaled Nova 1.0 Model (~191M Parameters)
    config = Nova1Config(
        vocab_size=tokenizer.vocab_size,
        d_model=args.d_model,
        n_heads=12,
        d_ff=3072,
        n_layers_L=4,
        n_layers_H=4,
        max_seq_len=512,
        learning_rate=args.lr,
        pad_token_id=tokenizer.pad_id,
        unk_token_id=tokenizer.unk_id,
        bos_token_id=tokenizer.bos_id,
        eos_token_id=tokenizer.eos_id,
        mask_token_id=tokenizer.mask_id
    )

    print(f"\n=======================================================", flush=True)
    print(f"🌟 Building Gemini-Style Nova 1.0 Engine", flush=True)
    print(f"Device: {config.device.upper()} | Precision: {config.dtype} | Vocab: {config.vocab_size:,}", flush=True)
    print(f"=======================================================\n", flush=True)

    # 3. Create DataLoader
    split_name = "train_sft" if "ultrachat" in args.dataset else "train"
    dataloader = create_hf_dataloader(
        dataset_name=args.dataset,
        subset=args.subset,
        tokenizer=tokenizer,
        max_seq_len=config.max_seq_len,
        batch_size=args.batch_size,
        split=split_name,
        max_samples=args.max_samples
    )

    # 4. Instantiate Model & Trainer
    model = Nova1HRM(config)
    num_params = sum(p.numel() for p in model.parameters())
    print(f"🚀 Model Total Trainable Parameters: {num_params:,} (~{num_params / 1e6:.1f} Million)", flush=True)

    trainer = Nova1Trainer(model=model, config=config, dataloader=dataloader)

    # 5. Train Model
    print(f"\nStarting Gemini-style pre-training on AMD ROCm GPU across {args.max_samples:,} samples...", flush=True)
    for epoch in range(args.epochs):
        avg_loss = trainer.train_epoch(epoch, args.epochs)
        print(f"Epoch {epoch+1}/{args.epochs} Complete — Average Loss: {avg_loss:.4f}", flush=True)

        # Save Checkpoint
        ckpt_path = f"checkpoints/nova1_gemini_epoch_{epoch+1}.pt"
        trainer.save_checkpoint(ckpt_path)

    final_path = "checkpoints/nova1_final.pt"
    trainer.save_checkpoint(final_path)
    print(f"\nGemini-style Nova 1.0 training complete! Model saved to '{final_path}'.", flush=True)


if __name__ == "__main__":
    main()