#!/usr/bin/env python3
"""
Fork-test pilot 2: same protocol as run_pilot.py but sampling from a
coherent completion model (gpt-3.5-turbo-instruct) instead of GPT-2-small,
and adding a first-sentence embedding variant (the sense-committing region).

Pilot 1 finding (kept, continuations_*.jsonl + results.json): GPT-2-small at
temp 1.0 is noise-dominated; no condition shows fork structure; sense signal
visible only in margin means (RIV < AMB < FIN). Instrument requires a model
whose continuations commit to a reading.
"""

import json
import time
from pathlib import Path

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

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


def generate(client):
    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"p2_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 first_sentence(text):
    # purely syntactic segmentation (allowed: format task, not semantics)
    t = text.strip()
    for stop in [". ", ".\n", "! ", "? "]:
        i = t.find(stop)
        if 0 < i < len(t) - 1:
            return t[:i + 1]
    return t


def embed_all(client, all_conts):
    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_full = {c: embed_batch(v) for c, v in all_conts.items()}
    embs_fs = {c: embed_batch([first_sentence(t) for t in v])
               for c, v in all_conts.items()}
    anchors = embed_batch([FIN_ANCHOR, RIV_ANCHOR])
    np.savez_compressed(
        OUT / "p2_embeddings.npz", anchors=anchors,
        **{f"full_{c}": e for c, e in embs_full.items()},
        **{f"fs_{c}": e for c, e in embs_fs.items()})
    return embs_full, embs_fs, anchors


def analyze(embs, anchors, tag):
    from sklearn.mixture import GaussianMixture
    fin_a, riv_a = normalize(anchors)
    eps_grid = np.arange(0.10, 0.85, 0.025)
    out = {}
    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(g2.means_.ravel()[0] - g2.means_.ravel()[1]) /
                    np.sqrt(g2.covariances_.ravel().mean()))
        out[cond] = {
            "margin_mean": float(margins.mean()),
            "margin_std": float(margins.std()),
            "delta_bic_k2_minus_k1": bic[2] - bic[1],
            "gmm2_component_means": [float(x) for x in g2.means_.ravel()],
            "gmm2_weights": [float(w) for w in g2.weights_],
            "gmm2_separation": sep,
            "pi0_sweep": pi0_sweep(X, eps_grid),
            "margins": [float(x) for x in margins],
        }
        d = out[cond]["delta_bic_k2_minus_k1"]
        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 out[cond]["pi0_sweep"] if 0.3 <= s["eps"] <= 0.65}
        print(f"[{tag}/{cond}] mean={margins.mean():+.3f} dBIC={d:+.1f} "
              f"sep={sep:.2f} -> {verdict}")
        print("   pi0:", " ".join(f"{e:.2f}:{c}" for e, c in comps.items()),
              flush=True)
    return out


if __name__ == "__main__":
    from openai import OpenAI
    client = OpenAI(api_key=load_env_key())
    conts = generate(client)
    embs_full, embs_fs, anchors = embed_all(client, conts)
    results = {
        "model": MODEL, "n_samples": N_SAMPLES, "max_tokens": MAX_TOKENS,
        "prompts": PROMPTS,
        "full": analyze(embs_full, anchors, "full"),
        "first_sentence": analyze(embs_fs, anchors, "fs"),
    }
    with open(OUT / "p2_results.json", "w") as f:
        json.dump(results, f, indent=2)
    print("done")
