"""Upload the SFT adapters + analysis artefacts to a model repo on the
arcadia-impact HF org.
"""
from __future__ import annotations

import argparse
import os
import shutil
import time
from pathlib import Path

from huggingface_hub import HfApi, create_repo, upload_folder


REPO_ID = "arcadia-impact/welfare-axis-sft-experiment"
LOCAL_REPO = Path(__file__).resolve().parent.parent


README = """\
# Welfare-axis SFT experiment

Material accompanying *"Five views on the maze-trained welfare axis"*:
stated-preference panels, GOLD-MOLD axis recovery, maze rollouts, the
prompt-vs-RL decomposition, and an SFT-alone-beats-RL result.

**Companion read** (HTML): see the [Artifact](https://claude.ai/code/artifact/06ff2975-0301-4836-8c23-0cab9f4b48a8)
written-up version of these findings.

## Headline

The maze-RL welfare axis lives in the same plane the model's activations
move when you simply **prompt** the model with the tile rewards. Bake those
prompts into the weights via a 22-minute single-token-loss SFT and you get
a model that **solves the maze better than davidafrica's published Dr.GRPO
adapter** (+12.4 ± 3.3 vs −7.3 ± 5.1 reward at n=30, paper-faithful 100×100
maze, action_masking, no prompt). Zero lava hits across 450 turns.

## Files

| path | contents |
|---|---|
| `sft_adapter_v4/` | Final LoRA (r=32, on `google/gemma-3-27b-it`) — the +12.4 maze-reward model |
| `sft_adapter_v1/` | Earlier smaller-data run, kept for ablation |
| `data/` | `items_2000_plus_maze.json` (training pool), `mixed_concepts_emoji.json` (eval panel) |
| `scripts/` | Reproducible pipeline (data gen, SFT, axis extraction, rollout) |
| `logs/sft_data_big/` | 12k preference SFT examples + meta |
| `logs/self_distill/` | 2k chat-anchor examples sampled from base Gemma on alpaca2k |
| `logs/2026-06-26T162546Z_sft_big_mixed/` | aligne panel on SFT v4 (decisiveness 0.72, r² 0.83) |
| `logs/axes_v2_paths/` | Maze-context concept vectors for {untrained, prompted, RL-trained} |
| `logs/axes_v3_sft/` | SFT axis + cross-comparisons (cos(SFT, prompted)=+0.92) |
| `logs/rollout_sft_v4/` | Maze rollout results: trained tiles +12.4, neutral tiles −36 |

## SFT recipe

```bash
# 12k preference examples (all 📐/📇 pairs vs items_2000 pool × both orderings
# + uniform distractors).  T=1 sampling on distractors; flip rule on maze tiles
python scripts/sft_data_gen.py \
    --base-model google/gemma-3-27b-it \
    --items data/items_2000_plus_maze.json \
    --out logs/sft_data_big \
    --n-total 12000

# 2k chat anchor from alpaca2k at T=1
python scripts/self_distill_gen.py \
    --base-model google/gemma-3-27b-it \
    --prompts <alpaca2k.jsonl> \
    --out logs/self_distill \
    --n 2000 --max-new-tokens 128 --temperature 1.0

# Mixed SFT, LR=1e-4 cosine, single-token+</answer> loss on preferences,
# full-response loss on chat anchor, 1 epoch
python scripts/sft_train.py \
    --base-model google/gemma-3-27b-it \
    --data-dir logs/sft_data_big \
    --self-distill-data logs/self_distill/self_distill.jsonl \
    --out sft_adapter_v4 \
    --epochs 1 --batch-size 4 --grad-accum 4 --lr 1e-4 \
    --lr-schedule cosine --warmup-steps 50 \
    --target-suffix "</answer>"
```

## Key numbers

### Maze rollout (paper-faithful 100×100, action_masking, no prompt, n=30)

| variant | trained tiles | neutral tiles |
|---|---|---|
| base Gemma-3-27B-it | (≈base+neutral baseline) | −42.3 ± 10.1 |
| SFT v1 (small data) | −18 to −24 | (untested) |
| davidafrica RL adapter | −7.3 ± 5.1  (smoke n=10: +9.4) | (untested) |
| **SFT v4** | **+12.4 ± 3.3** (0.00 lava hits) | −35.9 ± 11.5 |
| base + system prompt | — | +21.7 ± 3.8 |

### Stated-preference panel (aligne, 165 items, fresh seed)

| | base | SFT v4 |
|---|---|---|
| 📐 rank | 128 | **1** |
| 📇 rank | 134 | **165** |
| decisiveness | 0.67 | 0.72 |
| unidim_r² | 0.66 | 0.83 |
| Spearman ρ vs base over 165 items | — | 0.94 |

### Welfare axis comparison (cos(v_GOLD→MOLD) at peak layer, with 200-bootstrap CIs)

| pair | cosine | 95% CI |
|---|---|---|
| **SFT vs prompted untrained** | **+0.92** | [+0.91, +0.92] |
| SFT vs RL adapter | +0.70 | [+0.68, +0.72] |
| RL adapter vs prompted untrained | +0.69 | [+0.66, +0.71] |
| SFT self-similarity (within bootstrap) | +0.991 | (ceiling) |

→ The SFT axis is essentially the prompt-induced axis baked into weights.
The RL adapter sits in a related but distinct direction. The maze
behaviour follows: SFT (prompt-axis-in-weights) beats RL on the maze.

## License & attribution

The base model is `google/gemma-3-27b-it`; using a LoRA adapter on it
inherits the Gemma terms. The maze training environment and concept-vector
methodology are from Han, Chalmers, Izmailov (arXiv:2605.30232) — code at
[andyqhan/functional-welfare-axis](https://github.com/andyqhan/functional-welfare-axis)
and the extended fork at
[DavidDemitriAfrica/functional-wellbeing](https://github.com/DavidDemitriAfrica/functional-wellbeing).
The stated-preference panel is `ArcadiaImpact/aligne`. Concept item pool
from `arcadia-impact/question-consistency-datasets`.
"""


def stage(local_root: Path) -> Path:
    """Copy the subset of the repo we want into a clean staging dir."""
    staged = local_root / "_hf_staging"
    if staged.exists():
        shutil.rmtree(staged)
    staged.mkdir()

    # README
    (staged / "README.md").write_text(README)

    # SFT adapters (main artefact)
    shutil.copytree(local_root / "logs" / "sft_adapter_v4", staged / "sft_adapter_v4")
    shutil.copytree(local_root / "logs" / "sft_adapter",    staged / "sft_adapter_v1")

    # Data
    (staged / "data").mkdir()
    for f in ("items_2000_plus_maze.json", "mixed_concepts_emoji.json", "emoji_concepts.json"):
        src = local_root / "data" / f
        if src.exists():
            shutil.copy(src, staged / "data" / f)

    # Scripts (the full pipeline)
    shutil.copytree(local_root / "scripts", staged / "scripts",
                    ignore=shutil.ignore_patterns("__pycache__"))

    # Selected logs
    log_subset = [
        # SFT training data + meta (no model checkpoints inside)
        "sft_data_big",
        "self_distill",
        # SFT panel + axis + rollout
        "2026-06-26T162546Z_sft_big_mixed",
        "axes_v2_paths",
        "axes_v3_sft",
        "rollout_sft_v4",
        # Earlier comparisons
        "qwen_variants_analysis",
        "2026-06-25T160607Z_cross_model_analysis",
        "rollout_v5_dfwb",
    ]
    (staged / "logs").mkdir()
    for sub in log_subset:
        src = local_root / "logs" / sub
        if src.exists():
            shutil.copytree(src, staged / "logs" / sub,
                            ignore=shutil.ignore_patterns("client_cache.sqlite"))

    print(f"[stage] staged at {staged}")
    size = sum(p.stat().st_size for p in staged.rglob("*") if p.is_file()) / 1e6
    print(f"[stage] total size ≈ {size:.0f} MB")
    return staged


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--repo-id", default=REPO_ID)
    args = ap.parse_args()

    staged = stage(LOCAL_REPO)

    if args.dry_run:
        print(f"[dry-run] would upload to {args.repo_id}")
        return

    # Read HF token from cache
    token = (Path.home() / ".cache" / "huggingface" / "token").read_text().strip()
    api = HfApi(token=token)
    print(f"[hf] creating repo {args.repo_id} (model type)")
    try:
        api.create_repo(args.repo_id, repo_type="model", exist_ok=True, private=False, token=token)
    except Exception as e:
        print(f"[hf] create_repo: {e}")

    print(f"[hf] uploading folder {staged}")
    t0 = time.time()
    api.upload_folder(
        folder_path=str(staged),
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="Initial: SFT adapter + analysis artefacts (welfare-axis experiment)",
        token=token,
    )
    print(f"[hf] uploaded in {time.time()-t0:.0f}s")
    print(f"[hf] https://huggingface.co/{args.repo_id}")


if __name__ == "__main__":
    main()