#!/usr/bin/env python3
"""
split_automatability_decompose.py

Tests whether automatability benefits from a 2-stage decomposition:

  Stage 1 (binary): low_entropy {deterministic, template} vs needs_routing {specialist, slm, llm, human}
  Stage 2 (4-way):  specialist / slm / llm / human  (only for items routed past stage 1)

The flat 6-way model confuses deterministic↔template because the text features
that distinguish them aren't present in taskitem titles. This decomposition
isolates that gap — stage 1 correctly routes low-entropy items out, stage 2
handles the routed set cleanly.

Outputs combined accuracy as a 5-way result (low_entropy, specialist, slm, llm, human)
where deterministic/template are merged into one bucket, plus routing accuracy alone.

Usage:
    python3 specialists/scripts/split_automatability_decompose.py \\
        --data-dir  specialists/data/generated \\
        --runs-dir  /mnt/storage/fabricate_runs \\
        --results-out specialists/eval/results/automatability_decompose.json
"""

import argparse
import json
import subprocess
import sys
from collections import Counter
from pathlib import Path


# Must be at module level so torch.load can find __main__.Vocab when unpickling checkpoints
class Vocab:
    PAD, UNK = 0, 1

    def __init__(self):
        self.str_to_id = {"<pad>": 0, "<unk>": 1}
        self.id_to_str = ["<pad>", "<unk>"]

    def encode(self, text, seq_len):
        ids = [self.str_to_id.get(t, self.UNK) for t in text.lower().split()]
        ids = ids[:seq_len]
        ids += [self.PAD] * (seq_len - len(ids))
        return ids

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


ROOT         = Path(__file__).parent.parent.parent
VENV_PYTHON  = sys.executable
TRAIN_SCRIPT = ROOT / "specialists" / "scripts" / "train_specialist_pt.py"
EVAL_SCRIPT  = ROOT / "specialists" / "eval" / "eval_specialist_pt.py"

# Original 6-way label order: deterministic(0), template(1), specialist(2), slm(3), llm(4), human(5)
LABEL_NAMES_6 = ["deterministic", "template", "specialist", "slm", "llm", "human"]

# Stage 1: low_entropy = {0, 1}, needs_routing = {2, 3, 4, 5}
LOW_ENTROPY = {0, 1}

# Stage 2: remap {2→0, 3→1, 4→2, 5→3}
STAGE2_LABELS = ["specialist", "slm", "llm", "human"]
STAGE2_REMAP  = {2: 0, 3: 1, 4: 2, 5: 3}


def split_binary(src: Path, out: Path):
    """6-way TSV → binary TSV: low_entropy=0, needs_routing=1."""
    rows = []
    with open(src) as f:
        for line in f:
            parts = line.strip().split("\t", 2)
            if len(parts) < 3:
                continue
            orig = int(parts[0])
            rows.append(f"{'0' if orig in LOW_ENTROPY else '1'}\t0\t{parts[2]}")
    out.write_text("\n".join(rows) + "\n")
    n0 = sum(1 for r in rows if r.startswith("0"))
    print(f"  {out.name}: {len(rows)} rows  low_entropy={n0}  needs_routing={len(rows)-n0}")


def split_stage2(src: Path, out: Path):
    """6-way TSV → 4-way TSV keeping only needs_routing rows, remapped to 0-3."""
    rows = []
    with open(src) as f:
        for line in f:
            parts = line.strip().split("\t", 2)
            if len(parts) < 3:
                continue
            orig = int(parts[0])
            if orig in STAGE2_REMAP:
                rows.append(f"{STAGE2_REMAP[orig]}\t0\t{parts[2]}")
    out.write_text("\n".join(rows) + "\n")
    from collections import Counter
    dist = Counter(int(r.split("\t")[0]) for r in rows)
    print(f"  {out.name}: {len(rows)} rows  " + "  ".join(f"{STAGE2_LABELS[i]}={dist[i]}" for i in range(4)))


def train(run_dir: Path, train_tsv: Path, eval_tsv: Path, labels: str, max_steps: int = 10000) -> bool:
    run_dir.mkdir(parents=True, exist_ok=True)
    for f in run_dir.glob("*"):
        f.unlink()
    r = subprocess.run([
        str(VENV_PYTHON), str(TRAIN_SCRIPT),
        "--dataset",      str(train_tsv),
        "--eval-dataset", str(eval_tsv),
        "--out-dir",      str(run_dir),
        "--labels",       labels,
        "--lr",           "0.001",
        "--weight-decay", "0.01",
        "--max-steps",    str(max_steps),
        "--hidden-dim",   "256",
        "--layers",       "2",
        "--heads",        "4",
        "--batch-size",   "256",
        "--grok-loss-threshold", "0.05",
        "--grok-acc-jump",       "10.0",
        "--stop-after-grokking-blocks", "4",
        "--reset",
        "--history-out",  str(run_dir / "history.json"),
    ])
    return r.returncode == 0


def eval_checkpoint(run_dir: Path, eval_tsv: Path, labels: str) -> dict | None:
    ck = run_dir / "checkpoint.pt"
    if not ck.exists():
        return None
    r = subprocess.run(
        [str(VENV_PYTHON), str(EVAL_SCRIPT),
         "--checkpoint", str(ck),
         "--dataset", str(eval_tsv), "--labels", labels, "--quiet"],
        capture_output=True, text=True,
    )
    if r.returncode != 0:
        return None
    try:
        return json.loads(r.stdout.strip().splitlines()[-1])
    except (json.JSONDecodeError, IndexError):
        return None


def evaluate_combined(eval_tsv_6way: Path, stage1_run: Path, stage2_run: Path) -> dict:
    """Load both models; stage1 routes, stage2 classifies routed items.
    Reports: routing accuracy, 5-way accuracy (low_entropy merged), and
    comparison against original 6-way labels."""
    import torch
    import torch.nn as nn
    import numpy as np

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def _tok(text):
        return text.lower().split()

    class _Model(nn.Module):
        def __init__(self, cfg):
            super().__init__()
            D, L, H, V, seq = cfg["hidden_dim"], cfg["layers"], cfg["heads"], cfg["vocab_size"], cfg["seq_len"]
            self.tok_emb = nn.Embedding(V, D, padding_idx=0)
            self.pos_emb = nn.Embedding(seq, D)
            enc = nn.TransformerEncoderLayer(d_model=D, nhead=H, dim_feedforward=D*4,
                                             dropout=0.0, batch_first=True)
            self.encoder = nn.TransformerEncoder(enc, num_layers=L, enable_nested_tensor=False)
            self.head = nn.Linear(D, cfg["n_labels"])
            self.seq_len = seq

        def forward(self, x):
            pad_mask = (x == 0)
            pos = torch.arange(x.size(1), device=x.device).unsqueeze(0)
            h = self.tok_emb(x) + self.pos_emb(pos)
            h = self.encoder(h, src_key_padding_mask=pad_mask)
            lengths = (~pad_mask).float().sum(1, keepdim=True).clamp(min=1)
            h = (h * (~pad_mask).unsqueeze(-1).float()).sum(1) / lengths
            return self.head(h)

    def load_model(run_dir):
        ckpt = torch.load(run_dir / "checkpoint.pt", map_location=device, weights_only=False)
        m = _Model(ckpt["config"]).to(device)
        m.load_state_dict(ckpt["model"])
        m.eval()
        vocab = ckpt.get("vocab_obj") or ckpt.get("vocab")
        if isinstance(vocab, dict):
            class _V:
                UNK = 1
                def __init__(self, d): self.str_to_id = d
                def encode(self, text, seq_len):
                    ids = [self.str_to_id.get(t, self.UNK) for t in text.lower().split()]
                    ids = ids[:seq_len]; ids += [0] * (seq_len - len(ids)); return ids
            vocab = _V(vocab)
        return m, vocab, ckpt["config"]["seq_len"]

    def predict_probs(model, vocab, seq_len, texts):
        import torch.nn.functional as F
        all_probs = []
        with torch.no_grad():
            for start in range(0, len(texts), 256):
                batch = texts[start:start+256]
                ids = [vocab.encode(t, seq_len) for t in batch]
                x = torch.tensor(ids, dtype=torch.long, device=device)
                logits = model(x)
                all_probs.append(F.softmax(logits, dim=-1).cpu().numpy())
        return np.concatenate(all_probs, axis=0)

    # Load original 6-way eval data
    labels_gt, texts = [], []
    with open(eval_tsv_6way) as f:
        for line in f:
            parts = line.strip().split("\t", 2)
            if len(parts) < 3:
                continue
            labels_gt.append(int(parts[0]))
            texts.append(parts[2])
    labels_gt = np.array(labels_gt)
    n = len(labels_gt)

    m1, v1, s1 = load_model(stage1_run)
    m2, v2, s2 = load_model(stage2_run)

    probs1 = predict_probs(m1, v1, s1, texts)   # shape (n, 2): [low_entropy, needs_routing]
    probs2 = predict_probs(m2, v2, s2, texts)   # shape (n, 4): [specialist, slm, llm, human]

    stage1_pred = np.argmax(probs1, axis=1)  # 0=low_entropy, 1=needs_routing

    # Ground-truth routing: 0 if original label in {0,1}, else 1
    gt_routing = np.where(np.isin(labels_gt, [0, 1]), 0, 1)
    routing_acc = float((stage1_pred == gt_routing).mean() * 100)

    # Build 5-way predictions (low_entropy=0, specialist=1, slm=2, llm=3, human=4)
    stage2_pred = np.argmax(probs2, axis=1)  # 0-3 → specialist/slm/llm/human
    pred_5way = np.where(stage1_pred == 0, 0, stage2_pred + 1)

    # Ground-truth 5-way: low_entropy if orig in {0,1}, else stage2 remapped
    gt_5way = np.where(np.isin(labels_gt, [0, 1]), 0, labels_gt - 1)

    acc_5way = float((pred_5way == gt_5way).mean() * 100)
    correct_5way = pred_5way == gt_5way

    label_names_5 = ["low_entropy", "specialist", "slm", "llm", "human"]
    per_class = {}
    for i, name in enumerate(label_names_5):
        mask = gt_5way == i
        per_class[name] = {
            "acc": float(correct_5way[mask].mean() * 100) if mask.sum() > 0 else None,
            "n": int(mask.sum()),
        }

    errors = Counter()
    for gt, pred in zip(gt_5way.tolist(), pred_5way.tolist()):
        if gt != pred:
            errors[(label_names_5[gt], label_names_5[pred])] += 1

    return {
        "method": "2stage_decomposed",
        "routing_accuracy": round(routing_acc, 1),
        "accuracy_5way": round(acc_5way, 1),
        "n_eval": n,
        "per_class_5way": per_class,
        "top_errors": [{"gt": k[0], "pred": k[1], "count": v}
                       for k, v in errors.most_common(6)],
        "note": "low_entropy bucket = {deterministic, template} merged; 5-way is the effective label space",
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir",    default="specialists/data/generated")
    ap.add_argument("--runs-dir",    default="/mnt/storage/fabricate_runs")
    ap.add_argument("--results-out", default="specialists/eval/results/automatability_decompose.json")
    ap.add_argument("--skip-train",  action="store_true")
    args = ap.parse_args()

    data_dir  = Path(args.data_dir)
    runs_dir  = Path(args.runs_dir)
    results   = Path(args.results_out)
    results.parent.mkdir(parents=True, exist_ok=True)

    tmp = data_dir  # write split TSVs alongside existing data

    stage1_train = tmp / "automatability_binary_train.tsv"
    stage1_eval  = tmp / "automatability_binary_eval.tsv"
    stage2_train = tmp / "automatability_routed_train.tsv"
    stage2_eval  = tmp / "automatability_routed_eval.tsv"

    print("=== Splitting data ===")
    split_binary(data_dir / "automatability_train.tsv", stage1_train)
    split_binary(data_dir / "automatability_eval.tsv",  stage1_eval)
    split_stage2(data_dir / "automatability_train.tsv", stage2_train)
    split_stage2(data_dir / "automatability_eval.tsv",  stage2_eval)

    stage1_run = runs_dir / "whetstone_automatability_stage1_binary"
    stage2_run = runs_dir / "whetstone_automatability_stage2_4way"

    if not args.skip_train:
        print("\n=== Training stage 1: low_entropy vs needs_routing (binary) ===")
        train(stage1_run, stage1_train, stage1_eval, "low_entropy,needs_routing")

        print("\n=== Training stage 2: specialist/slm/llm/human (4-way) ===")
        train(stage2_run, stage2_train, stage2_eval, "specialist,slm,llm,human")

    print("\n=== Evaluating individual models ===")
    r1 = eval_checkpoint(stage1_run, stage1_eval, "low_entropy,needs_routing")
    r2 = eval_checkpoint(stage2_run, stage2_eval, "specialist,slm,llm,human")
    if r1:
        print(f"  Stage 1 (binary):  {r1.get('overall_accuracy')}%")
    if r2:
        print(f"  Stage 2 (4-way):   {r2.get('overall_accuracy')}%")

    print("\n=== Evaluating combined 2-stage pipeline ===")
    combined = evaluate_combined(data_dir / "automatability_eval.tsv", stage1_run, stage2_run)
    print(f"  Routing accuracy:  {combined['routing_accuracy']}%")
    print(f"  5-way accuracy:    {combined['accuracy_5way']}%")
    print("  Per class (5-way):")
    for name, stat in combined["per_class_5way"].items():
        if stat["acc"] is not None:
            print(f"    {name:20s}: {stat['acc']:.1f}%  (n={stat['n']})")

    output = {
        "gate": "automatability",
        "flat_6way_baseline": {"small_plus": 69.6, "medium": 78.3},
        "decomposed": {
            "stage1_binary": r1,
            "stage2_4way":   r2,
            "combined":      combined,
        },
    }
    results.write_text(json.dumps(output, indent=2))
    print(f"\nResults → {results}")


if __name__ == "__main__":
    main()