#!/usr/bin/env python3
"""
Fork-test pilot: behavioral pi_0 estimation of filler spaces (TTS paper, §6).

Three horns (prompts ending at the sign 'bank'):
  AMB — ambiguous boundary (both senses live)        -> predict 2 stable components
  FIN — financially disambiguated boundary           -> predict 1 component
  RIV — riparian disambiguated boundary              -> predict 1 component

Protocol per horn:
  1. Sample N continuations from GPT-2 (pure ancestral sampling: temp=1.0,
     no top-k/top-p truncation) — unbiased draws from the model's filler space.
  2. Embed continuation text with text-embedding-3-small.
  3. pi_0 at resolution eps := connected components of the eps-graph
     (= single-linkage clusters cut at cosine distance eps), counting only
     components with >= MIN_COMP members (dust ignored).
  4. Sense margin per draw: cos(e, fin_anchor) - cos(e, riv_anchor);
     GMM BIC k=1 vs k=2 on margins as the bimodality check.

Outputs: results.json, embeddings.npz, continuations JSONL per condition.
"""

import json
import os
import sys
import time
from pathlib import Path

import numpy as np

OUT = Path(__file__).parent
N_SAMPLES = 200
MAX_NEW_TOKENS = 35
BATCH = 20
MIN_COMP = 5
SEED = 11235

PROMPTS = {
    "AMB": "He stood up, grabbed his coat, and walked down to the bank.",
    "FIN": "He needed to deposit the check before noon, so he walked down to the bank.",
    "RIV": "He carried his fishing rod through the reeds and walked down to the bank.",
}

FIN_ANCHOR = ("The bank as a financial institution: money, accounts, deposits, "
              "tellers, loans, cash withdrawals, branch opening hours.")
RIV_ANCHOR = ("The bank as the side of a river: water, shore, reeds, mud, "
              "fishing, current, the edge of the stream.")


def load_env_key():
    for line in open("/home/iman/cassie-project/.env"):
        line = line.strip()
        if line.startswith("OPENAI_API_KEY="):
            return line.split("=", 1)[1].strip().strip('"').strip("'")
    raise RuntimeError("OPENAI_API_KEY not found in .env")


def generate():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch.manual_seed(SEED)
    tok = AutoTokenizer.from_pretrained("gpt2")
    tok.pad_token = tok.eos_token
    tok.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained("gpt2")
    model.eval()

    all_conts = {}
    for cond, prompt in PROMPTS.items():
        conts = []
        t0 = time.time()
        enc = tok([prompt] * BATCH, return_tensors="pt", padding=True)
        while len(conts) < N_SAMPLES:
            with torch.no_grad():
                out = model.generate(
                    **enc,
                    do_sample=True,
                    temperature=1.0,
                    top_k=0,           # no truncation: ancestral sampling
                    top_p=1.0,
                    max_new_tokens=MAX_NEW_TOKENS,
                    pad_token_id=tok.eos_token_id,
                )
            new = out[:, enc["input_ids"].shape[1]:]
            conts.extend(tok.batch_decode(new, skip_special_tokens=True))
            print(f"[{cond}] {len(conts)}/{N_SAMPLES} ({time.time()-t0:.0f}s)",
                  flush=True)
        all_conts[cond] = conts[:N_SAMPLES]
        with open(OUT / f"continuations_{cond}.jsonl", "w") as f:
            for c in all_conts[cond]:
                f.write(json.dumps({"prompt": prompt, "continuation": c}) + "\n")
    return all_conts


def embed(all_conts):
    from openai import OpenAI
    client = OpenAI(api_key=load_env_key())

    def embed_batch(texts):
        vecs = []
        for i in range(0, len(texts), 100):
            resp = client.embeddings.create(
                model="text-embedding-3-small",
                input=[t if t.strip() else " " for t in texts[i:i+100]],
            )
            vecs.extend([d.embedding for d in resp.data])
        return np.array(vecs)

    embs = {cond: embed_batch(conts) for cond, conts in all_conts.items()}
    anchors = embed_batch([FIN_ANCHOR, RIV_ANCHOR])
    np.savez_compressed(OUT / "embeddings.npz",
                        anchors=anchors,
                        **{f"emb_{c}": e for c, e in embs.items()})
    return embs, anchors


def normalize(X):
    return X / np.linalg.norm(X, axis=1, keepdims=True)


def pi0_sweep(X, eps_grid):
    """Connected components of the eps-graph = single-linkage cut at eps."""
    from scipy.cluster.hierarchy import fcluster, linkage
    from scipy.spatial.distance import pdist
    D = pdist(normalize(X), metric="cosine")
    Z = linkage(D, method="single")
    out = []
    for eps in eps_grid:
        labels = fcluster(Z, t=eps, criterion="distance")
        sizes = np.bincount(labels)[1:]
        out.append({
            "eps": float(eps),
            "components": int((sizes > 0).sum()),
            "components_min%d" % MIN_COMP: int((sizes >= MIN_COMP).sum()),
            "largest": int(sizes.max()),
        })
    return out


def analyze(embs, anchors):
    from sklearn.mixture import GaussianMixture

    fin_a, riv_a = normalize(anchors)
    eps_grid = np.arange(0.10, 0.85, 0.025)
    results = {"n_samples": N_SAMPLES, "model": "gpt2",
               "max_new_tokens": MAX_NEW_TOKENS, "seed": SEED,
               "prompts": PROMPTS, "conditions": {}}

    for cond, X in embs.items():
        Xn = normalize(X)
        margins = Xn @ fin_a - Xn @ riv_a
        gmm = {}
        for k in (1, 2):
            g = GaussianMixture(k, n_init=10, random_state=SEED)
            g.fit(margins.reshape(-1, 1))
            gmm[f"bic_k{k}"] = float(g.bic(margins.reshape(-1, 1)))
        g2 = GaussianMixture(2, n_init=10, random_state=SEED)
        lab = g2.fit_predict(margins.reshape(-1, 1))
        means = g2.means_.ravel()
        sep = float(abs(means[0] - means[1]) /
                    np.sqrt(g2.covariances_.ravel().mean()))
        results["conditions"][cond] = {
            "margin_mean": float(margins.mean()),
            "margin_std": float(margins.std()),
            "gmm": gmm,
            "delta_bic_k2_minus_k1": gmm["bic_k2"] - gmm["bic_k1"],
            "gmm2_component_means": [float(m) for m in means],
            "gmm2_weights": [float(w) for w in g2.weights_],
            "gmm2_separation": sep,
            "pi0_sweep": pi0_sweep(X, eps_grid),
            "margins": [float(m) for m in margins],
            "gmm2_labels": [int(l) for l in lab],
        }

    with open(OUT / "results.json", "w") as f:
        json.dump(results, f, indent=2)
    return results


def report(results):
    print("\n========== FORK PILOT SUMMARY ==========")
    for cond, r in results["conditions"].items():
        d = r["delta_bic_k2_minus_k1"]
        verdict = "BIMODAL (fork signature)" if d < -10 else \
                  "unimodal (canonical signature)" if d > -2 else "marginal"
        print(f"\n[{cond}] margin mean={r['margin_mean']:+.3f} "
              f"std={r['margin_std']:.3f}")
        print(f"  GMM dBIC(k2-k1) = {d:+.1f}  -> {verdict}")
        print(f"  GMM2 means={['%+.3f' % m for m in r['gmm2_component_means']]} "
              f"weights={['%.2f' % w for w in r['gmm2_weights']]} "
              f"sep={r['gmm2_separation']:.2f}")
        rows = [s for s in r["pi0_sweep"] if 0.3 <= s["eps"] <= 0.65]
        comp = {s["eps"]: s[f"components_min{MIN_COMP}"] for s in rows}
        print(f"  pi0(eps) [min size {MIN_COMP}]: "
              + " ".join(f"{e:.2f}:{c}" for e, c in comp.items()))


if __name__ == "__main__":
    if "--analyze-only" in sys.argv:
        data = np.load(OUT / "embeddings.npz")
        embs = {c: data[f"emb_{c}"] for c in PROMPTS}
        anchors = data["anchors"]
    else:
        conts = generate()
        embs, anchors = embed(conts)
    results = analyze(embs, anchors)
    report(results)
