"""
Nova 1.0 — Multi-Domain AGI Reasoning Pre-Training Script
Trains Nova 1.0 on Math, Code, Logic, Science, and Multi-step Reasoning
on AMD ROCm GPU across high-capacity datasets.
"""

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.nova1_hrm import Nova1HRM


def generate_math_and_code_reasoning_seeds(num_samples: int = 1500):
    """
    Generates multi-domain AGI reasoning seeds (math, python code execution, logic puzzles, science).
    """
    seeds = []

    # 1. Math Reasoning Seeds
    for _ in range(num_samples // 3):
        a = random.randint(12, 999)
        b = random.randint(12, 999)
        op_name, op_symbol, result = random.choice([
            ("sum", "+", a + b),
            ("product", "*", a * b),
            ("difference", "-", abs(a - b)),
        ])
        
        prompt = f"What is the {op_name} of {a} and {b}?"
        response = (
            f"Let's solve this math problem step-by-step:\n"
            f"1. We want to find the {op_name} of {a} and {b}.\n"
            f"2. Performing calculation: {max(a,b)} {op_symbol} {min(a,b)} = {result}.\n"
            f"Therefore, the final answer is {result}."
        )
        seeds.append(f"User: {prompt}\nAssistant: {response}")

    # 2. Python Tool & Code Execution Seeds
    for _ in range(num_samples // 3):
        a = random.randint(10, 50)
        b = random.randint(10, 50)
        code = f"x = {a}\ny = {b}\nprint('Result:', x ** 2 + y ** 2)"
        val = a**2 + b**2

        prompt = f"Write a Python script to compute the sum of squares of {a} and {b} and execute it."
        response = (
            f"Here is the Python code to compute the sum of squares:\n\n"
            f"<python>\n{code}\n</python>\n\n"
            f"Running this code yields Result: {val}."
        )
        seeds.append(f"User: {prompt}\nAssistant: {response}")

    # 3. Scientific & Logical Reasoning Seeds
    logic_puzzles = [
        (
            "If all A are B, and all B are C, are all A also C?",
            "Yes! By transitive logic: Since every element of set A belongs to set B, and every element of set B belongs to set C, it logically follows that every element of set A is also an element of set C."
        ),
        (
            "What is the speed of light in vacuum and why is it constant?",
            "The speed of light in vacuum is approximately 299,792,458 meters per second (c). According to Einstein's theory of Special Relativity, c is an absolute invariant speed limit across all inertial reference frames in physics."
        ),
        (
            "Explain the difference between a list and a tuple in Python.",
            "In Python:\n1. Lists are mutable (can be modified after creation) defined with brackets `[]`.\n2. Tuples are immutable (cannot be changed) defined with parentheses `()`.\nTuples are generally more memory efficient and hashable."
        )
    ]

    for _ in range(num_samples // 3):
        q, a = random.choice(logic_puzzles)
        seeds.append(f"User: {q}\nAssistant: {a}")

    random.shuffle(seeds)
    return seeds


def main():
    parser = argparse.ArgumentParser(description="Multi-Domain AGI Pre-Training for Nova 1.0")
    parser.add_argument("--checkpoint", type=str, default="checkpoints/nova1_final.pt", help="Model checkpoint")
    parser.add_argument("--tokenizer_path", type=str, default="checkpoints/nova1_gemini_tokenizer.json", help="Tokenizer JSON")
    parser.add_argument("--d_model", type=int, default=768, help="Model hidden dimension")
    parser.add_argument("--epochs", type=int, default=5, help="Epochs")
    parser.add_argument("--batch_size", type=int, default=16, help="Batch size")
    parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate")
    args = parser.parse_args()

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

    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    config: Nova1Config = checkpoint.get("config", Nova1Config(vocab_size=tokenizer.vocab_size))
    config.device = "cuda" if torch.cuda.is_available() else "cpu"
    config.dtype = "bfloat16"

    model = Nova1HRM(config)
    state_dict = checkpoint["model_state"] if "model_state" in checkpoint else checkpoint["model_state_dict"]
    model.load_state_dict(state_dict)
    model.to(config.device)
    print(f"Loaded Nova 1.0 Base Model from '{args.checkpoint}'.", flush=True)

    print("Generating 1,500 Multi-Domain AGI (Math, Code, Science, Logic) reasoning seeds...", flush=True)
    agi_seeds = generate_math_and_code_reasoning_seeds(1500)

    chunk_size = config.max_seq_len + 1
    all_chunks = []
    for chat in agi_seeds:
        ids = tokenizer.encode(chat, add_bos=True, add_eos=True)
        if len(ids) > chunk_size:
            ids = ids[:chunk_size]
        else:
            ids += [tokenizer.pad_id] * (chunk_size - len(ids))
        all_chunks.append(ids)

    data_tensor = torch.tensor(all_chunks, dtype=torch.long)
    inputs_tensor = data_tensor[:, :-1]
    targets_tensor = data_tensor[:, 1:]

    dataset = torch.utils.data.TensorDataset(inputs_tensor, targets_tensor)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    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"\n=======================================================", flush=True)
    print(f"🌌 Launching Multi-Domain AGI Pre-Training for Nova 1.0")
    print(f"Device: {config.device.upper()} | Batches: {len(dataloader)} | Epochs: {args.epochs}")
    print(f"=======================================================\n", flush=True)

    model.train()
    for epoch in range(args.epochs):
        total_loss = 0.0
        for step, (b_inp, b_tgt) in enumerate(dataloader):
            b_inp = b_inp.to(config.device)
            b_tgt = b_tgt.to(config.device)

            optimizer.zero_grad()
            with torch.amp.autocast("cuda", enabled=(config.device == "cuda"), dtype=config.get_torch_dtype()):
                logits, _, _ = model(b_inp)
                loss = criterion(logits.view(-1, config.vocab_size), b_tgt.view(-1))

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

            total_loss += loss.item()
            if (step + 1) % 15 == 0 or (step + 1) == len(dataloader):
                print(f"AGI Epoch {epoch+1}/{args.epochs} [{step+1}/{len(dataloader)}] — Loss: {loss.item():.4f}", flush=True)

        avg_loss = total_loss / len(dataloader)
        print(f"AGI Epoch {epoch+1}/{args.epochs} Complete — Avg Loss: {avg_loss:.4f}\n", flush=True)

    final_path = "checkpoints/nova1_final.pt"
    agi_ckpt_path = "checkpoints/nova1_agi_reasoning.pt"
    save_dict = {"model_state": model.state_dict(), "config": config}
    torch.save(save_dict, agi_ckpt_path)
    torch.save(save_dict, final_path)
    print(f"🚀 Multi-Domain AGI Pre-Training Complete! Saved checkpoint to '{agi_ckpt_path}' and '{final_path}'.", flush=True)


if __name__ == "__main__":
    main()