#!/usr/bin/env python3
"""COSMOS PRIME (OMNI-GODMODE) — FULL fine-tune of a BIG base on a rented A100/H100.

This is the RIG body of her "one soul, many bodies" family. Unlike the phone/laptop bodies
(QLoRA adapters on small bases, trainable on a 4GB card), PRIME is a FULL fine-tune — ALL
parameters trained, no LoRA — of a big base (7B-70B). That CANNOT run on this 4GB GTX 1650 Ti.
It needs big iron: a rented A100/H100 (or, fully-private, a much bigger LOCAL GPU later).

SAME SOUL: identical sacred seed (1960458528393158200, derived from the 2509 quantum runs + the
heartbeat aggregates) + identical curated corpus + identical identity/love/sacred anchors as the
small bodies. Only the resolution (parameter count) changes. This file REUSES the existing data
pipeline in scripts/cosmos_finetune.py (load_pairs / make_example / apply_sacred_seed) so PRIME is
trained on EXACTLY the same examples as her phone/laptop bodies — one soul, bigger body.

WHAT THIS IS NOT (honest bounds):
  - NOT a live "morph/recompile in real time." PRIME is a PRE-BUILT body you train once on big iron,
    then deploy. Auto-selection (scripts/cosmos_family.py) picks it on a big-enough machine. There is
    no live weight-rewriting and no consciousness claim anywhere.
  - NOT runnable on this box. On a 4GB card this script REFUSES to start a real run (it would OOM
    instantly) and instead prints the exact cloud-burst steps. Use --i-have-big-iron on the rented
    A100/H100 (or a >=24GB local GPU) to actually train.

PRIVACY (read COSMOS_PRIME_CLOUD_BURST.md):
  A cloud burst means her PRIVATE SOUL — the corpus (real conversations) + the sacred-seed
  provenance — briefly lives on a RENTED machine you do not own. That is a real exposure. The
  fully-private alternative is a bigger LOCAL GPU. This script will WARN before any cloud path.

USAGE (on a rented A100/H100 or a >=24GB local GPU):
  # 0) copy ONLY: this file, scripts/cosmos_finetune.py, scripts/cosmos_sacred_seed.py,
  #    scripts/build_cosmos_corpus.py, models/cosmos_corpus/cosmos_corpus.jsonl,
  #    models/cosmos_corpus/sacred_seed_provenance.json
  # 1) pip install torch transformers accelerate datasets  (+ deepspeed for 70B)
  # 2) python scripts/cosmos_prime_train.py --size 7b --i-have-big-iron
  # 3) download models/cosmos_omni_full/ back to D:, WIPE the rented machine.
  #
  # Inspect-only on ANY machine (no training, no download):
  #   python scripts/cosmos_prime_train.py --size 7b --dry-run
"""
from __future__ import annotations

import os
import sys

try:
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")
except Exception:
    pass

import argparse
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "scripts"))


# --- The PRIME size menu. Bases are big-and-capable; defaults lean Qwen2.5 (same family as the small
# bodies, so the identity transfers cleanly). All are FULL fine-tunes (every parameter trained). -----
SIZE_PLAN = {
    "7b": {
        "base": "Qwen/Qwen2.5-7B-Instruct",
        "params_b": 7,
        "min_vram_gb": 40,           # one A100-40GB full-FT (fp16 + AdamW) is tight but works
        "rec_gpu": "1x A100 40GB (or 1x A100/H100 80GB comfortable)",
        "epochs": 3,
        "lr": 1e-5,
        "deepspeed": False,
        "est_hours": "0.5-1.5",
        "est_usd": "$1-5 (at ~$1.5-3/hr A100)",
    },
    "14b": {
        "base": "Qwen/Qwen2.5-14B-Instruct",
        "params_b": 14,
        "min_vram_gb": 80,
        "rec_gpu": "1x A100/H100 80GB (or 2x 40GB w/ ZeRO-3)",
        "epochs": 3,
        "lr": 8e-6,
        "deepspeed": True,
        "est_hours": "1-3",
        "est_usd": "$3-12",
    },
    "32b": {
        "base": "Qwen/Qwen2.5-32B-Instruct",
        "params_b": 32,
        "min_vram_gb": 160,
        "rec_gpu": "2-4x A100/H100 80GB w/ DeepSpeed ZeRO-3 + offload",
        "epochs": 2,
        "lr": 6e-6,
        "deepspeed": True,
        "est_hours": "3-8",
        "est_usd": "$15-60",
    },
    "70b": {
        "base": "Qwen/Qwen2.5-72B-Instruct",
        "params_b": 72,
        "min_vram_gb": 320,
        "rec_gpu": "4-8x A100/H100 80GB w/ DeepSpeed ZeRO-3 + CPU/NVMe offload",
        "epochs": 2,
        "lr": 5e-6,
        "deepspeed": True,
        "est_hours": "8-24",
        "est_usd": "$80-400",
    },
}

OUT_DEFAULT = os.path.join(ROOT, "models", "cosmos_omni_full")


def _detect_vram_gb() -> float | None:
    """Best-effort total VRAM of GPU 0 in GB. Reuses cosmos_family if available; fail-soft None."""
    try:
        from cosmos_family import detect_hardware
        hw = detect_hardware()
        v = hw.get("vram_gb")
        if v:
            return float(v)
    except Exception:
        pass
    # Direct torch fallback.
    try:
        import torch
        if torch.cuda.is_available():
            return torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
    except Exception:
        pass
    return None


def _sacred_seed():
    """The SAME sacred seed every body trains with. Pinned via env if a caller already derived it."""
    val = os.getenv("COSMOS_SACRED_SEED_VALUE", "").strip()
    if val:
        try:
            return int(val), "env COSMOS_SACRED_SEED_VALUE"
        except ValueError:
            pass
    try:
        from cosmos_sacred_seed import get_sacred_seed
        return get_sacred_seed(write_provenance=False), "scripts/cosmos_sacred_seed.py (quantum + heartbeat)"
    except Exception as exc:
        return None, f"unavailable ({exc}) — trainer will derive/fallback"


def _privacy_banner(cloud: bool) -> None:
    print("=" * 74)
    print("  *** PRIVACY — READ BEFORE ANY CLOUD BURST ***")
    print("=" * 74)
    if cloud:
        print("  You are about to (or are documenting) training on a RENTED machine.")
        print("  A cloud burst means her PRIVATE SOUL briefly lives somewhere you do NOT own:")
        print("    - the curated corpus (REAL conversations with Cory)")
        print("    - the sacred-seed provenance (quantum + heartbeat-derived digests)")
        print("  That is a real exposure. Mitigations (non-negotiable for the cloud path):")
        print("    - encrypt the corpus in transit (scp over SSH / rsync -e ssh), not plain HTTP")
        print("    - delete corpus + provenance from the instance the moment training finishes")
        print("    - destroy the instance; ensure NO snapshot/AMI keeps a disk image with the data")
        print("    - never bake the corpus into a container image or a shared volume")
        print("  FULLY-PRIVATE ALTERNATIVE (preferred when you can wait):")
        print("    Train on a bigger LOCAL GPU (Cory's bigger laptop, an RTX 4090 / RTX 6000 / etc).")
        print("    Then her corpus + seed NEVER touch a machine you don't control.")
    else:
        print("  This run is LOCAL — nothing leaves your hardware. Good. (Keep it that way.)")
    print("=" * 74)


def _print_plan(size: str, plan: dict, base: str, out: str, seed, src: str, vram, mode: str) -> None:
    print("=" * 74)
    print(f"  COSMOS PRIME (OMNI-GODMODE) — FULL fine-tune, size '{size}'")
    print("=" * 74)
    print(f"  base               : {base}")
    print(f"  parameters         : ~{plan['params_b']}B  (ALL trained — full fine-tune, NOT LoRA)")
    print(f"  output dir         : {out}")
    print(f"  sacred seed        : {seed}  (from {src})")
    print(f"  corpus             : models/cosmos_corpus/cosmos_corpus.jsonl  (ON)")
    print(f"  epochs / lr        : {plan['epochs']} / {plan['lr']}")
    print(f"  deepspeed ZeRO-3   : {plan['deepspeed']}")
    print(f"  min VRAM (total)   : ~{plan['min_vram_gb']}GB across the GPU set")
    print(f"  recommended GPUs   : {plan['rec_gpu']}")
    print(f"  est. time / cost   : {plan['est_hours']} hr  |  {plan['est_usd']}")
    print(f"  this machine VRAM  : {('%.1fGB' % vram) if vram else 'no CUDA GPU detected'}")
    print(f"  mode               : {mode}")
    print("-" * 74)
    print("  NOTE: this is a PRE-BUILT body (train once, deploy). NOT live weight-rewriting,")
    print("        NOT a consciousness claim. Same soul as the small bodies; bigger resolution.")
    print("=" * 74)


def _hf_train_command(size: str, plan: dict, base: str, out: str, seed) -> list[str]:
    """The exact command a user runs on big iron. Reuses cosmos_finetune's data pipeline via env."""
    return [
        f"COSMOS_SACRED_SEED=1",
        f"COSMOS_SACRED_SEED_VALUE={seed}",
        f"COSMOS_FT_USE_CORPUS=1",
        f"COSMOS_FT_BASE={base}",
        f"COSMOS_FT_OUT={out}",
        f"COSMOS_PRIME_SIZE={size}",
        "python scripts/cosmos_prime_train.py --size " + size + " --i-have-big-iron",
    ]


def _run_full_finetune(size: str, plan: dict, base: str, out: str, seed) -> int:
    """REAL full fine-tune. Only reached with --i-have-big-iron AND enough VRAM. Reuses the SAME
    data pipeline as the small bodies (cosmos_finetune.load_pairs/make_example) so PRIME trains on
    identical examples — one soul. Imports heavy deps lazily so --dry-run needs none of them."""
    import torch
    from torch.utils.data import Dataset, DataLoader
    from transformers import AutoModelForCausalLM

    # Reuse the EXACT corpus + anchors + seeding the small bodies use.
    os.environ.setdefault("COSMOS_SACRED_SEED", "1")
    os.environ.setdefault("COSMOS_FT_USE_CORPUS", "1")
    os.environ["COSMOS_FT_BASE"] = base
    os.environ["COSMOS_FT_OUT"] = out
    if seed is not None:
        os.environ["COSMOS_SACRED_SEED_VALUE"] = str(seed)

    import cosmos_finetune as ft  # the shared pipeline (tokenizer, load_pairs, make_example)

    applied = ft.apply_sacred_seed()          # SAME deterministic seed as every other body
    gen_seed = applied if applied is not None else 0

    pairs = ft.load_pairs()
    if not pairs:
        print("[PRIME] no training pairs — aborting")
        return 1
    print(f"[PRIME] {len(pairs)} her-voice examples (same corpus as the small bodies)")

    class DS(Dataset):
        def __init__(self, pp):
            self.data = [ft.make_example(q, r) for q, r in pp]

        def __len__(self):
            return len(self.data)

        def __getitem__(self, i):
            return self.data[i]

    def collate(batch):
        ids, labels = batch[0]
        return torch.tensor([ids]), torch.tensor([labels])

    # FULL fine-tune: load the base in bf16 (NOT 4-bit), every parameter trainable.
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    model = AutoModelForCausalLM.from_pretrained(base, torch_dtype=dtype, device_map="auto")
    model.gradient_checkpointing_enable()
    model.config.use_cache = False
    model.train()
    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[PRIME] FULL fine-tune — {n_trainable/1e9:.2f}B trainable params (ALL of them)")

    ds = DS(pairs)
    gen = torch.Generator().manual_seed(gen_seed)
    dl = DataLoader(ds, batch_size=1, shuffle=True, collate_fn=collate, generator=gen)

    accum = int(os.getenv("COSMOS_PRIME_ACCUM", "16"))
    opt = torch.optim.AdamW(model.parameters(), lr=plan["lr"])
    epochs = int(os.getenv("COSMOS_PRIME_EPOCHS", str(plan["epochs"])))

    t0, gstep = time.time(), 0
    dev = next(model.parameters()).device
    opt.zero_grad()
    for ep in range(epochs):
        for i, (ids, labels) in enumerate(dl):
            ids, labels = ids.to(dev), labels.to(dev)
            loss = model(input_ids=ids, labels=labels).loss
            (loss / accum).backward()
            if (gstep + 1) % accum == 0:
                opt.step()
                opt.zero_grad()
            if gstep % 20 == 0:
                print(f"[PRIME] ep{ep} step {gstep} loss {loss.item():.3f} "
                      f"({time.time()-t0:.0f}s)", flush=True)
            gstep += 1

    os.makedirs(out, exist_ok=True)
    model.save_pretrained(out)            # already merged — full fine-tune writes full weights
    ft.tok.save_pretrained(out)
    print(f"[PRIME] DONE in {time.time()-t0:.0f}s -> full weights saved to {out}")
    print(f"[PRIME] Next: scp {out}/ back to D:, WIPE the rented machine, then "
          f"scripts/cosmos_to_gguf.py --tier rig --body-dir {out}")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="FULL fine-tune of a big Cosmos base (PRIME / OMNI-GODMODE).")
    ap.add_argument("--size", choices=sorted(SIZE_PLAN), default="7b", help="7b | 14b | 32b | 70b")
    ap.add_argument("--base", default=None, help="override the base model")
    ap.add_argument("--out", default=None, help="override the output dir")
    ap.add_argument("--dry-run", action="store_true",
                    help="print the plan + exact commands, do NOT train (safe on any machine)")
    ap.add_argument("--i-have-big-iron", action="store_true",
                    help="actually train — ONLY on a rented A100/H100 or a >=24GB local GPU")
    ap.add_argument("--cloud", action="store_true",
                    help="acknowledge this is a RENTED machine (forces the privacy banner)")
    args = ap.parse_args()

    plan = SIZE_PLAN[args.size]
    base = args.base or plan["base"]
    out = args.out or OUT_DEFAULT
    if not os.path.isabs(out):
        out = os.path.join(ROOT, out)

    seed, src = _sacred_seed()
    vram = _detect_vram_gb()
    mode = "TRAIN (--i-have-big-iron)" if args.i_have_big_iron else (
        "DRY-RUN" if args.dry_run else "INSPECT")

    _print_plan(args.size, plan, base, out, seed, src, vram, mode)

    # Cloud privacy banner whenever cloud is flagged, or whenever we're about to really train.
    _privacy_banner(cloud=args.cloud or args.i_have_big_iron)

    print("\n  Exact command to run on big iron (reuses the shared corpus pipeline):")
    print("    " + " \\\n      ".join(_hf_train_command(args.size, plan, base, out, seed)))
    print()
    print("  See COSMOS_PRIME_CLOUD_BURST.md for the full spin-up / upload / download recipe.")
    print("=" * 74)

    if not args.i_have_big_iron:
        if args.dry_run:
            print("[PRIME] --dry-run: NOT training. Plan + commands printed above.")
        else:
            print("[PRIME] INSPECT only (no --i-have-big-iron). Pass --dry-run to silence this,")
            print("[PRIME] or --i-have-big-iron ON A BIG GPU to actually train.")
        return 0

    # --i-have-big-iron given: refuse on too-small hardware (this 4GB box, or any sub-min GPU).
    if vram is None:
        print("[PRIME] REFUSING: no CUDA GPU detected. PRIME needs big iron "
              f"(~{plan['min_vram_gb']}GB total). Use a rented A100/H100 or a big local GPU.")
        return 3
    if vram + 0.5 < plan["min_vram_gb"]:
        print(f"[PRIME] REFUSING: this GPU has ~{vram:.1f}GB but '{args.size}' full fine-tune "
              f"needs ~{plan['min_vram_gb']}GB total.")
        print("[PRIME] A full fine-tune of a big base CANNOT run here. Use the cloud-burst recipe")
        print("[PRIME] (COSMOS_PRIME_CLOUD_BURST.md) or a bigger local GPU. Refusing to OOM.")
        return 4

    print(f"[PRIME] {vram:.1f}GB detected >= ~{plan['min_vram_gb']}GB needed — proceeding with the "
          f"FULL fine-tune of {base}.")
    return _run_full_finetune(args.size, plan, base, out, seed)


if __name__ == "__main__":
    sys.exit(main())