#!/usr/bin/env python3
"""
Fork-test pilot 3: calibrated balanced boundary (C5 from calibrate.py),
n=300 per condition, gpt-3.5-turbo-instruct, first-sentence embeddings.
Adds 2-means separation statistic in full embedding space.
"""

import json
import time
from pathlib import Path

import numpy as np
from run_pilot import (FIN_ANCHOR, RIV_ANCHOR, MIN_COMP, load_env_key,
                       normalize, pi0_sweep)
from run_pilot2 import first_sentence

OUT = Path(__file__).parent
N_SAMPLES = 300
MAX_TOKENS = 40
MODEL = "gpt-3.5-turbo-instruct"
SEED = 11235

PROMPTS = {
    "AMB": "We met them near the bank, just before sunset.",          # calibrated 22/28
    "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.",
}


def main():
    from openai import OpenAI
    from sklearn.cluster import KMeans
    from sklearn.metrics import silhouette_score
    from sklearn.mixture import GaussianMixture

    client = OpenAI(api_key=load_env_key())

    def embed(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)

    all_conts = {}
    for cond, prompt in PROMPTS.items():
        conts = []
        while len(conts) < N_SAMPLES:
            n = min(50, N_SAMPLES - len(conts))
            resp = client.completions.create(
                model=MODEL, prompt=prompt, max_tokens=MAX_TOKENS,
                temperature=1.0, top_p=1.0, n=n)
            conts.extend(c.text for c in resp.choices)
            print(f"[{cond}] {len(conts)}/{N_SAMPLES}", flush=True)
            time.sleep(1)
        all_conts[cond] = conts[:N_SAMPLES]
        with open(OUT / f"p3_continuations_{cond}.jsonl", "w") as f:
            for c in all_conts[cond]:
                f.write(json.dumps({"prompt": prompt, "continuation": c}) + "\n")

    embs = {c: embed([first_sentence(t) for t in v])
            for c, v in all_conts.items()}
    anchors = embed([FIN_ANCHOR, RIV_ANCHOR])
    np.savez_compressed(OUT / "p3_embeddings.npz", anchors=anchors,
                        **{f"fs_{c}": e for c, e in embs.items()})

    fin_a, riv_a = normalize(anchors)
    eps_grid = np.arange(0.10, 0.85, 0.025)
    results = {"model": MODEL, "n_samples": N_SAMPLES, "prompts": PROMPTS,
               "conditions": {}}
    for cond, X in embs.items():
        Xn = normalize(X)
        margins = Xn @ fin_a - Xn @ riv_a
        m = margins.reshape(-1, 1)
        bic = {k: float(GaussianMixture(k, n_init=10, random_state=SEED)
                        .fit(m).bic(m)) for k in (1, 2)}
        g2 = GaussianMixture(2, n_init=10, random_state=SEED).fit(m)
        sep = float(abs(np.diff(g2.means_.ravel())[0]) /
                    np.sqrt(g2.covariances_.ravel().mean()))
        km = KMeans(2, n_init=10, random_state=SEED).fit(Xn)
        sil = float(silhouette_score(Xn, km.labels_, metric="cosine"))
        results["conditions"][cond] = {
            "margin_mean": float(margins.mean()),
            "margin_std": float(margins.std()),
            "delta_bic_k2_minus_k1": bic[2] - bic[1],
            "gmm2_means": [float(x) for x in g2.means_.ravel()],
            "gmm2_weights": [float(w) for w in g2.weights_],
            "gmm2_separation": sep,
            "kmeans2_silhouette_cosine": sil,
            "kmeans2_sizes": [int(s) for s in np.bincount(km.labels_)],
            "pi0_sweep": pi0_sweep(X, eps_grid),
            "margins": [float(x) for x in margins],
        }
        d = bic[2] - bic[1]
        verdict = ("BIMODAL (fork)" if d < -10 else
                   "unimodal (canonical)" if d > -2 else "marginal")
        comps = {s["eps"]: s[f"components_min{MIN_COMP}"]
                 for s in results["conditions"][cond]["pi0_sweep"]
                 if 0.25 <= s["eps"] <= 0.65}
        print(f"\n[{cond}] mean={margins.mean():+.3f} dBIC={d:+.1f} "
              f"sep={sep:.2f} sil2={sil:.3f} "
              f"sizes={np.bincount(km.labels_).tolist()} -> {verdict}")
        print("   pi0:", " ".join(f"{e:.2f}:{c}" for e, c in comps.items()),
              flush=True)

    json.dump(results, open(OUT / "p3_results.json", "w"), indent=2)
    print("\ndone")


if __name__ == "__main__":
    main()
