"""
Script 2: Pre-train Scaled Nova 1.0 Model from Scratch (High Capacity)
Leverages full system resources: 32 CPU Threads, 119GB RAM, AMD ROCm.
"""

import os
import argparse
import torch

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 Scaled Nova 1.0 Model")
    parser.add_argument("--dataset", type=str, default="tatsu-lab/alpaca", help="HuggingFace dataset name")
    parser.add_argument("--subset", type=str, default=None, help="HuggingFace dataset subset/config name")
    parser.add_argument("--tokenizer_path", type=str, default="checkpoints/nova1_tokenizer.json", help="Tokenizer path")
    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="Directory to save model checkpoints")
    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("--d_model", type=int, default=768, help="Model hidden dimension")
    parser.add_argument("--max_samples", type=int, default=25000, help="Number of HuggingFace dataset samples")
    args = parser.parse_args()

    # 1. Load tokenizer
    if not os.path.exists(args.tokenizer_path):
        print(f"Tokenizer not found at {args.tokenizer_path}. Please run scripts/train_tokenizer.py first!")
        return

    tokenizer = Nova1Tokenizer.load(args.tokenizer_path)
    print(f"Loaded subword tokenizer with vocab size {tokenizer.vocab_size}.")

    # 2. Build Nova 1.0 Config (Scaled)
    config = Nova1Config(
        vocab_size=tokenizer.vocab_size,
        d_model=args.d_model,
        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"Initialized Scaled Nova 1.0 Config: Device={config.device}, dtype={config.dtype}, d_model={config.d_model}, max_seq_len={config.max_seq_len}")

    # 3. Create DataLoader from Hugging Face dataset
    print(f"Preparing high-capacity DataLoader for dataset '{args.dataset}' ({args.max_samples:,} samples)...")
    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,
        max_samples=args.max_samples
    )

    # 4. Instantiate Nova 1.0 Model & Trainer
    model = Nova1HRM(config)
    num_params = sum(p.numel() for p in model.parameters())
    print(f"\n=======================================================")
    print(f"🚀 Scaled Nova 1.0 Model Architecture Built!")
    print(f"Total Trainable Parameters: {num_params:,} (~{num_params / 1e6:.1f} Million)")
    print(f"=======================================================\n")

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

    # 5. Run Training Loop
    print(f"Starting Nova 1.0 Pre-training from scratch using {config.num_threads} CPU threads & {config.device.upper()}...")
    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}")

        # Save checkpoint after each epoch
        ckpt_path = os.path.join(args.checkpoint_dir, f"nova1_epoch_{epoch+1}.pt")
        trainer.save_checkpoint(ckpt_path)

    # Save final model checkpoint
    final_path = os.path.join(args.checkpoint_dir, "nova1_final.pt")
    trainer.save_checkpoint(final_path)
    print(f"\nScaled Nova 1.0 training complete! Final model saved to '{final_path}'.")


if __name__ == "__main__":
    main()