#!/usr/bin/env python3
"""
Convert HRM checkpoints to safetensors and push to Hugging Face.

Usage:
    python push_to_hf.py \
        --baseline "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 belligerent-squirrel/step_52080" \
        --tiered "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HRM_Tiered realistic-dalmatian/step_52080" \
        --repo "Code2aum/HRM_based_evolution"
"""

import argparse
import os
import shutil
import torch
from safetensors.torch import save_file
from huggingface_hub import login, HfApi, upload_folder


def convert_checkpoint(ckpt_path, output_dir, model_name):
    """Load .pt checkpoint, save as safetensors + config."""
    os.makedirs(output_dir, exist_ok=True)

    # Load weights
    state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)

    # Strip torch.compile prefix if present
    cleaned = {}
    for k, v in state_dict.items():
        key = k.removeprefix("_orig_mod.")
        cleaned[key] = v

    # Save as safetensors
    st_path = os.path.join(output_dir, "model.safetensors")
    save_file(cleaned, st_path)
    print(f"  ✓ Saved {st_path} ({os.path.getsize(st_path)/1e6:.1f} MB)")

    # Copy config
    ckpt_dir = os.path.dirname(ckpt_path)
    config_src = os.path.join(ckpt_dir, "all_config.yaml")
    if os.path.exists(config_src):
        shutil.copy2(config_src, os.path.join(output_dir, "config.yaml"))
        print(f"  ✓ Copied config.yaml")

    # Copy model source
    for src_file in ["hrm_act_v1.py", "hrm_tiered.py", "losses.py"]:
        src_path = os.path.join(ckpt_dir, src_file)
        if os.path.exists(src_path):
            shutil.copy2(src_path, os.path.join(output_dir, src_file))

    # Write a model card
    card = f"""---
tags:
  - hrm
  - hierarchical-reasoning
  - sudoku
  - pytorch
license: mit
---

# {model_name}

Hierarchical Reasoning Model trained on Sudoku-Extreme-1K (20,000 epochs).

## Architecture
- **Type**: {model_name}
- **Hidden Size**: 512
- **Heads**: 8
- **H/L Layers**: 4/4
- **H/L Cycles**: 2/2
- **Parameters**: ~27.3M

## Training
- **Dataset**: Sudoku-Extreme-1K (1000 puzzles, 1000 augmentations each)
- **Epochs**: 20,000
- **Batch Size**: 384
- **Learning Rate**: 7e-5
- **GPU**: NVIDIA RTX 4090

## Usage
```python
from safetensors.torch import load_file
state_dict = load_file("{model_name}/model.safetensors")
```
"""
    with open(os.path.join(output_dir, "README.md"), "w") as f:
        f.write(card)
    print(f"  ✓ Created README.md")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--baseline", type=str, required=True)
    parser.add_argument("--tiered", type=str, required=True)
    parser.add_argument("--repo", type=str, default="Code2aum/HRM_based_evolution")
    parser.add_argument("--output-dir", type=str, default="hf_upload")
    parser.add_argument("--no-push", action="store_true", help="Convert only, skip push")
    args = parser.parse_args()

    print("=" * 60)
    print("  HRM → SafeTensors → Hugging Face")
    print("=" * 60)

    # Convert baseline
    base_dir = os.path.join(args.output_dir, "baseline_hrm_v1")
    print(f"\n  Converting Baseline...")
    convert_checkpoint(args.baseline, base_dir, "HRM_Baseline_V1")

    # Convert tiered
    tier_dir = os.path.join(args.output_dir, "tiered_hrm_sram_dram")
    print(f"\n  Converting Tiered...")
    convert_checkpoint(args.tiered, tier_dir, "HRM_Tiered_SRAM_DRAM")

    # Copy benchmark results if available
    bench_dir = "benchmark_results"
    if os.path.exists(bench_dir):
        dest = os.path.join(args.output_dir, "benchmark_results")
        if not os.path.exists(dest):
            shutil.copytree(bench_dir, dest)
            print(f"\n  ✓ Copied benchmark_results/")

    if args.no_push:
        print(f"\n  Files ready at: {args.output_dir}/")
        print("  Run without --no-push to upload to HF.")
        return

    # Push to HF
    print(f"\n  Logging in to Hugging Face...")
    login()

    print(f"\n  Uploading to {args.repo}...")
    api = HfApi()
    api.create_repo(repo_id=args.repo, repo_type="model", exist_ok=True)
    upload_folder(
        folder_path=args.output_dir,
        repo_id=args.repo,
        repo_type="model",
    )
    print(f"\n  ✓ Uploaded to https://huggingface.co/{args.repo}")

    print("\n" + "=" * 60)
    print("  Done!")
    print("=" * 60)


if __name__ == "__main__":
    main()