#!/usr/bin/env python3
"""
NVIDIA Nsight Systems Profiler for HRM Memory Tiering.

This script runs a few iterations of both the baseline and tiered models
and is designed to be executed via `nsys profile`.

Usage:
  nsys profile -t cuda,nvtx --stats=true --force-overwrite=true -o hrm_profile python run_nsys_profiler.py
"""

import torch
import torch.cuda.nvtx as nvtx
import argparse

from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
from models.hrm.hrm_tiered import HRM_Tiered
from models.memory_tier import MemoryTierManager
from run_training_comparison import DummyLossModel, create_dummy_batch

def profile_model(model_name, model, batch, iterations, device):
    print(f"Profiling {model_name}...")
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    model.train()

    # Warmup
    for _ in range(2):
        optimizer.zero_grad()
        carry = model.initial_carry(batch)
        carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
        carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
        carry.steps = carry.steps.to(device)
        carry.halted = carry.halted.to(device)
        carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}
        
        _, loss, _, _, _ = model(carry, batch, return_keys=[])
        loss.backward()
        optimizer.step()

    torch.cuda.synchronize()

    # Profiling Phase
    with torch.autograd.profiler.emit_nvtx():
        nvtx.range_push(f"{model_name}_Training_Loop")
        for i in range(iterations):
            nvtx.range_push(f"Iteration_{i}")
            optimizer.zero_grad()
            
            nvtx.range_push("Forward_Pass")
            carry = model.initial_carry(batch)
            carry.inner_carry.z_H = carry.inner_carry.z_H.to(device)
            carry.inner_carry.z_L = carry.inner_carry.z_L.to(device)
            carry.steps = carry.steps.to(device)
            carry.halted = carry.halted.to(device)
            carry.current_data = {k: v.to(device) for k, v in carry.current_data.items()}

            _, loss, _, _, _ = model(carry, batch, return_keys=[])
            nvtx.range_pop() # End Forward
            
            nvtx.range_push("Backward_Pass")
            loss.backward()
            optimizer.step()
            nvtx.range_pop() # End Backward
            
            nvtx.range_pop() # End Iteration
        nvtx.range_pop() # End Loop

    torch.cuda.synchronize()
    print(f"Finished {model_name}.\n")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch-size', type=int, default=16)
    parser.add_argument('--seq-len', type=int, default=128)
    parser.add_argument('--hidden-size', type=int, default=1024)
    parser.add_argument('--iterations', type=int, default=5)
    args = parser.parse_args()

    device = torch.device('cuda')
    vocab_size = 32

    config_dict = {
        'batch_size': args.batch_size,
        'seq_len': args.seq_len,
        'puzzle_emb_ndim': 0,
        'num_puzzle_identifiers': args.batch_size,
        'vocab_size': vocab_size,
        'H_cycles': 2,
        'L_cycles': 2,
        'H_layers': 4,
        'L_layers': 4,
        'hidden_size': args.hidden_size,
        'expansion': 4.0,
        'num_heads': 8,
        'pos_encodings': 'rope',
        'halt_max_steps': 1,
        'halt_exploration_prob': 0.0,
    }

    batch = create_dummy_batch(args.batch_size, args.seq_len, vocab_size, device)

    # 1. Baseline
    baseline = HierarchicalReasoningModel_ACTV1(config_dict).to(device)
    baseline_wrapped = DummyLossModel(baseline)
    profile_model("HRM_Baseline", baseline_wrapped, batch, args.iterations, device)
    del baseline_wrapped, baseline
    torch.cuda.empty_cache()

    # 2. Tiered
    mem_mgr = MemoryTierManager(device=device, enable_tracking=False)
    tiered = HRM_Tiered(config_dict, memory_manager=mem_mgr).to(device)
    tiered_wrapped = DummyLossModel(tiered)
    profile_model("HRM_Tiered", tiered_wrapped, batch, args.iterations, device)

if __name__ == "__main__":
    main()