#!/usr/bin/env python3
"""
2-horn (fill-in-the-middle) fork test — TTS paper §6.

The pilot (run_pilot3.py) sampled cones: prefix fixed, completion free.
This experiment holds BOTH flanking commitments fixed — the FIM setting —
and samples the bridge: a direct sample of Fill(σ) for the inner 2-horn.

Design: ONE ambiguous prefix (the calibrated bank boundary), THREE suffixes:
  NEU — neutral right context (does not disambiguate)   -> predict: fork preserved
  FIN — financially disambiguating right context        -> predict: fork killed (financial)
  RIV — riparian disambiguating right context           -> predict: fork killed (riparian)
Model: gpt-4o-mini chat, instruction-based infilling, temperature 1.0, n=200/condition.
Instrument: text-embedding-3-small on the whole infill; same margins + eps-sweep
as the pilot; plus cross-comparison to the pilot-3 cone distribution.
"""

import json
import time
from pathlib import Path

import numpy as np

import sys
sys.path.insert(0, str(Path(__file__).parent))
from run_pilot import FIN_ANCHOR, RIV_ANCHOR, MIN_COMP, load_env_key, normalize, pi0_sweep

OUT = Path(__file__).parent
N_SAMPLES = 200
MODEL = "gpt-3.5-turbo-instruct"   # prompted FIM on a completion model:
# instruction-tuned chat models (gpt-4o-mini tested) mode-collapse the
# bridge distribution and under-weight the suffix constraint; the
# completion model retains sampling diversity and respects both flanks.
SEED = 11235

PREFIX = "We met them near the bank, just before sunset."

SUFFIXES = {
    "NEU": "By the time we said our goodbyes, the stars were already out.",
    "FIN": "When the manager finally unlocked the doors again, we were the last customers of the day.",
    "RIV": "We watched the current carry our little boat downstream until it disappeared.",
}

FIM_TEMPLATE = (
    "Fill in the missing middle of the story so it reads as one coherent "
    "whole.\n\nSTORY BEFORE THE GAP: {p}\n\nSTORY AFTER THE GAP: {s}\n\n"
    "THE MISSING MIDDLE (one or two sentences):"
)


def main(n_samples=N_SAMPLES, conditions=None, tag="h", native=False):
    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())
    conds = conditions or list(SUFFIXES)

    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_infills = {}
    for cond in conds:
        infills = []
        prompt = FIM_TEMPLATE.format(p=PREFIX, s=SUFFIXES[cond])
        while len(infills) < n_samples:
            n = min(50, n_samples - len(infills))
            if native:
                resp = client.completions.create(
                    model=MODEL, temperature=1.0, top_p=1.0, n=n,
                    max_tokens=70, prompt=PREFIX + " ",
                    suffix=" " + SUFFIXES[cond])
            else:
                resp = client.completions.create(
                    model=MODEL, temperature=1.0, top_p=1.0, n=n,
                    max_tokens=70, prompt=prompt)
            infills.extend((c.text or "").strip() for c in resp.choices)
            print(f"[{cond}] {len(infills)}/{n_samples}", flush=True)
            time.sleep(0.5)
        all_infills[cond] = infills[:n_samples]
        with open(OUT / f"{tag}_infills_{cond}.jsonl", "w") as f:
            for c in all_infills[cond]:
                f.write(json.dumps({"prefix": PREFIX, "suffix": SUFFIXES[cond],
                                    "infill": c}) + "\n")

    embs = {c: embed(v) for c, v in all_infills.items()}
    anchors = embed([FIN_ANCHOR, RIV_ANCHOR])
    np.savez_compressed(OUT / f"{tag}_embeddings.npz", anchors=anchors,
                        **{f"in_{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, "prefix": PREFIX,
               "suffixes": SUFFIXES, "fim_template": FIM_TEMPLATE, "native_fim": native,
               "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,
            "pi0_sweep": pi0_sweep(X, eps_grid),
            "margins": [float(x) for x in margins],
        }
        d = bic[2] - bic[1]
        verdict = ("BIMODAL (fork signature)" if d < -10 else
                   "unimodal" 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} sil={sil:.3f} -> {verdict}")
        print("   pi0:", " ".join(f"{e:.2f}:{c}" for e, c in comps.items()),
              flush=True)

    json.dump(results, open(OUT / f"{tag}_results.json", "w"), indent=2)
    print("\ndone")


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--calibrate", action="store_true",
                    help="n=10 per condition, print samples")
    ap.add_argument("--native", action="store_true",
                    help="use the API's native suffix-param FIM (tag hn)")
    args = ap.parse_args()
    if args.native:
        main(tag="hn", native=True)
    elif args.calibrate:
        main(n_samples=10, tag="hcal")
        for cond in SUFFIXES:
            print(f"\n===== {cond} samples =====")
            for i, line in enumerate(open(OUT / f"hcal_infills_{cond}.jsonl")):
                if i >= 5: break
                print(" ", repr(json.loads(line)["infill"][:130]))
    else:
        main()
