#!/usr/bin/env python3
"""
Push all trained SAE and transcoder checkpoints to HuggingFace Hub.

Uses HUGGING_FACE_HUB_TOKEN from .env file.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parent.parent / ".env", override=False)

from huggingface_hub import HfApi, create_repo

CHECKPOINT_DIR = Path("circuit_tracer/data/checkpoints")
REPO_ID = "sarel/creditscope-circuit-models"

def main():
    token = os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
    if not token:
        print("ERROR: No HuggingFace token found. Set HUGGING_FACE_HUB_TOKEN or HF_TOKEN.")
        sys.exit(1)

    api = HfApi(token=token)

    # Create repo if it doesn't exist
    try:
        create_repo(REPO_ID, token=token, repo_type="model", exist_ok=True)
        print(f"Repo ready: {REPO_ID}")
    except Exception as e:
        print(f"Repo creation: {e}")

    # Upload all .pt checkpoints
    pt_files = sorted(CHECKPOINT_DIR.glob("*.pt"))
    print(f"Found {len(pt_files)} checkpoint files to upload")

    # Upload registry-compatible files (sae_lN.pt and tc_lN.pt)
    registry_files = [f for f in pt_files if f.name.startswith(("sae_l", "tc_l"))]
    print(f"Uploading {len(registry_files)} registry-compatible checkpoints...")

    for f in registry_files:
        size_mb = f.stat().st_size / 1024 / 1024
        print(f"  Uploading {f.name} ({size_mb:.1f} MB)...", end=" ", flush=True)
        try:
            api.upload_file(
                path_or_fileobj=str(f),
                path_in_repo=f"checkpoints/{f.name}",
                repo_id=REPO_ID,
                repo_type="model",
            )
            print("OK")
        except Exception as e:
            print(f"FAILED: {e}")

    # Also upload best/final variants
    extra_files = [f for f in pt_files if f not in registry_files]
    if extra_files:
        print(f"\nUploading {len(extra_files)} additional checkpoints (best/final)...")
        for f in extra_files:
            size_mb = f.stat().st_size / 1024 / 1024
            print(f"  Uploading {f.name} ({size_mb:.1f} MB)...", end=" ", flush=True)
            try:
                api.upload_file(
                    path_or_fileobj=str(f),
                    path_in_repo=f"checkpoints/{f.name}",
                    repo_id=REPO_ID,
                    repo_type="model",
                )
                print("OK")
            except Exception as e:
                print(f"FAILED: {e}")

    # Create a model card
    model_card = """---
tags:
  - sparse-autoencoder
  - transcoder
  - circuit-tracing
  - mechanistic-interpretability
  - qwen3.5
license: apache-2.0
---

# CreditScope Circuit Tracing Models

Sparse Autoencoders (SAEs) and MoE Transcoders trained on **Qwen3.5-35B-A3B-FP8** for mechanistic interpretability / circuit tracing.

## Models

### SAEs (Sparse Autoencoders)
- **Architecture**: JumpReLU, d_model=2048 → 16384 features (8x expansion)
- **Layers**: 0, 5, 10, 15, 20, 25, 30, 35, 39
- **Training**: 500 diverse prompts, ~5000 tokens, 2000-15000 steps per model
- **Files**: `sae_l{N}.pt`

### Transcoders (MoE Transcoders)
- **Architecture**: ReLU encoder/decoder, d_model=2048 → 16384 features
- **Layers**: 0, 5, 10, 15, 20, 25, 30, 35, 39
- **Training**: Maps pre-MoE residual to post-MoE output (learns MoE residual contribution)
- **Files**: `tc_l{N}.pt`

## Usage

```python
from circuit_tracer.saes.sparse_autoencoder import SparseAutoencoder
from circuit_tracer.transcoders.moe_transcoder import MoETranscoder

# Load SAE
sae = SparseAutoencoder.load("checkpoints/sae_l0.pt")

# Load transcoder
tc = MoETranscoder.load("checkpoints/tc_l0.pt")
```

## Training Details
- **Base model**: Qwen/Qwen3.5-35B-A3B-FP8
- **Activation collection**: Direct model forward hooks on 500 diverse prompts
- **SAE optimizer**: Adam, lr=3e-4, cosine annealing
- **TC optimizer**: Adam, lr=1e-3, cosine annealing
"""

    try:
        api.upload_file(
            path_or_fileobj=model_card.encode(),
            path_in_repo="README.md",
            repo_id=REPO_ID,
            repo_type="model",
        )
        print("\nModel card uploaded")
    except Exception as e:
        print(f"Model card upload failed: {e}")

    print(f"\nDone! View at: https://huggingface.co/{REPO_ID}")


if __name__ == "__main__":
    main()