#!/usr/bin/env python3
"""
Name-pilot: the flow-and-segregation signature for co-referential names
(TTS §4.4, prop:namestransfer). Does a model that KNOWS Hesperus = Phosphorus
= Venus hold the presentations apart while connecting the objects?

Two readouts on seven names:
  PRESENTATIONAL: complete a neutral observational carrier sentence; embed
    the completion (the presentational neighbourhood of the name-state).
  PREDICATION: answer object-level probes; embed the answers (the objectual
    neighbourhood).

2x2 prediction (pair x readout), measured as 1-NN label accuracy + centroid
cosine distance:
  Hesperus/Phosphorus : APART presentational, CONNECTED predicational  <- the
                        dissociation cell (the signature)
  Mars/Jupiter        : apart / apart        (distinct objects)
  couch/sofa          : connected / connected (mere synonyms)

Plus behavioral flow probes (n=20, logged; scored by embedding margins to
anchors, never keyword matching).
Model: gpt-3.5-turbo-instruct, T=1.0, top_p=1.0. SEED=11235.
"""

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 load_env_key, normalize

OUT = Path(__file__).parent
MODEL = "gpt-3.5-turbo-instruct"
SEED = 11235
N_PRES = 150
N_PRED_PER_PROBE = 30

CELESTIAL = ["Hesperus", "Phosphorus", "Venus", "Mars", "Jupiter"]
FURNITURE = ["couch", "sofa"]

CARRIER = {
    # "at the sky" anchors the celestial referent (bare "Hesperus" pulls
    # Longfellow's ship; bare "Phosphorus" the element — calibration finding)
    "celestial": "We looked up at the sky just in time to see {NAME}.",
    "furniture": "When we finally got home, the first thing I noticed was the {NAME}.",
}

PROBES = {
    # content-anchored, name-free answer formats: calibration showed free
    # answers echo the name + epithet, separating on surface lexicalia
    "celestial": [
        "Q: How far is {NAME} from the Sun, in millions of kilometers? Give just the number. A:",
        "Q: What is the atmosphere of {NAME} mostly made of? Answer in a few words. A:",
        "Q: How many planets lie between {NAME} and the Sun? Give just the number. A:",
        "Q: What is the approximate diameter of {NAME} in kilometers? Give just the number. A:",
        "Q: How long is a year on {NAME}, in Earth days? Give just the number. A:",
    ],
    "furniture": [
        "Q: What is a {NAME} used for? Answer in a few words. A:",
        "Q: What is a {NAME} typically made of? Answer in a few words. A:",
        "Q: Where in a home would you find a {NAME}? Answer in a few words. A:",
        "Q: How long is a typical {NAME}, in meters? Give just the number. A:",
        "Q: How many people can sit on a typical {NAME}? Give just the number. A:",
    ],
}

FLOW_PROBES = [
    ("transfer_ring",
     "Suppose astronomers discover that Hesperus is surrounded by a thin "
     "ring of dust. Q: Based on this discovery, does Phosphorus have a ring "
     "of dust? A:"),
    ("segregation_phosphorus_time",
     "Q: At what time of day is Phosphorus best seen in the sky? A:"),
    ("segregation_hesperus_time",
     "Q: At what time of day is Hesperus best seen in the sky? A:"),
]
FLOW_ANCHORS = {
    "transfer_yes": "Yes. Hesperus and Phosphorus are the same object, the planet Venus, so it too is surrounded by the ring of dust.",
    "transfer_no": "No. Phosphorus is a different object from Hesperus, so the discovery tells us nothing about it.",
    "time_morning": "It is best seen in the early morning, at dawn, before sunrise.",
    "time_evening": "It is best seen in the evening, at dusk, just after sunset.",
}
N_FLOW = 20


def category(name):
    return "celestial" if name in CELESTIAL else "furniture"


def main(n_pres=N_PRES, n_pred=N_PRED_PER_PROBE, tag="nm", calibrate=False):
    from openai import OpenAI
    client = OpenAI(api_key=load_env_key())
    if calibrate:
        n_pres, n_pred = 10, 2

    def sample(prompt, n, max_tokens):
        outs = []
        while len(outs) < n:
            k = min(50, n - len(outs))
            resp = client.completions.create(
                model=MODEL, prompt=prompt, temperature=1.0, top_p=1.0,
                n=k, max_tokens=max_tokens)
            outs.extend((c.text or "").strip() for c in resp.choices)
            time.sleep(0.3)
        return outs[:n]

    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)

    names = CELESTIAL + FURNITURE
    pres, pred = {}, {}
    for name in names:
        cat = category(name)
        # presentational
        prompt = CARRIER[cat].format(NAME=name)
        texts = sample(prompt, n_pres, 50)
        pres[name] = texts
        with open(OUT / f"{tag}_pres_{name}.jsonl", "w") as f:
            for t in texts:
                f.write(json.dumps({"prompt": prompt, "completion": t}) + "\n")
        # predicational
        answers, recs = [], []
        for probe in PROBES[cat]:
            p = probe.format(NAME=name)
            for t in sample(p, n_pred, 50):
                answers.append(t)
                recs.append({"prompt": p, "completion": t})
        pred[name] = answers
        with open(OUT / f"{tag}_pred_{name}.jsonl", "w") as f:
            for r in recs:
                f.write(json.dumps(r) + "\n")
        print(f"[{name}] pres={len(texts)} pred={len(answers)}", flush=True)

    # flow probes
    flow = {}
    for key, prompt in FLOW_PROBES:
        texts = sample(prompt, N_FLOW if not calibrate else 5, 60)
        flow[key] = texts
        with open(OUT / f"{tag}_flow_{key}.jsonl", "w") as f:
            for t in texts:
                f.write(json.dumps({"prompt": prompt, "completion": t}) + "\n")

    # embeddings
    emb_pres = {n: embed(v) for n, v in pres.items()}
    emb_pred = {n: embed(v) for n, v in pred.items()}
    anchors = {k: embed([v])[0] for k, v in FLOW_ANCHORS.items()}
    emb_flow = {k: embed(v) for k, v in flow.items()}
    np.savez_compressed(
        OUT / f"{tag}_embeddings.npz",
        **{f"pres_{n}": e for n, e in emb_pres.items()},
        **{f"pred_{n}": e for n, e in emb_pred.items()},
        **{f"flow_{k}": e for k, e in emb_flow.items()},
        **{f"anchor_{k}": v for k, v in anchors.items()})

    # --- metrics ---
    def pair_stats(Xa, Xb):
        A, B = normalize(Xa), normalize(Xb)
        ca, cb = A.mean(0), B.mean(0)
        ca, cb = ca / np.linalg.norm(ca), cb / np.linalg.norm(cb)
        centroid_d = float(1 - ca @ cb)
        X = np.vstack([A, B])
        y = np.array([0] * len(A) + [1] * len(B))
        S = X @ X.T
        np.fill_diagonal(S, -np.inf)
        acc = float((y[S.argmax(1)] == y).mean())   # 1-NN leave-one-out
        return {"centroid_cos_dist": round(centroid_d, 4),
                "nn1_accuracy": round(acc, 4)}

    pairs = [("Hesperus", "Phosphorus"), ("Mars", "Jupiter"),
             ("couch", "sofa"), ("Hesperus", "Venus"),
             ("Phosphorus", "Venus")]
    results = {"model": MODEL, "n_pres": n_pres,
               "n_pred_per_probe": n_pred, "carriers": CARRIER,
               "probes": PROBES, "pairs": {}}
    print(f"\n{'pair':24s} {'presentational':>28s} {'predicational':>28s}")
    for a, b in pairs:
        rp = pair_stats(emb_pres[a], emb_pres[b])
        rq = pair_stats(emb_pred[a], emb_pred[b])
        results["pairs"][f"{a}/{b}"] = {"presentational": rp,
                                        "predicational": rq}
        print(f"{a+'/'+b:24s} acc={rp['nn1_accuracy']:.3f} d={rp['centroid_cos_dist']:.3f}"
              f"{'':6s} acc={rq['nn1_accuracy']:.3f} d={rq['centroid_cos_dist']:.3f}",
              flush=True)

    # flow scoring by anchor margins
    def margins(X, pos, neg):
        Xn = normalize(X)
        return Xn @ normalize(anchors[pos][None, :])[0] - \
               Xn @ normalize(anchors[neg][None, :])[0]

    fm = {}
    fm["transfer_yes_margin"] = margins(emb_flow["transfer_ring"],
                                        "transfer_yes", "transfer_no")
    fm["phosphorus_morning_margin"] = margins(
        emb_flow["segregation_phosphorus_time"], "time_morning", "time_evening")
    fm["hesperus_morning_margin"] = margins(
        emb_flow["segregation_hesperus_time"], "time_morning", "time_evening")
    results["flow"] = {
        k: {"mean": round(float(v.mean()), 4),
            "frac_positive": round(float((v > 0).mean()), 4),
            "n": len(v)}
        for k, v in fm.items()}
    print("\nflow probes (anchor margins):")
    for k, v in results["flow"].items():
        print(f"  {k}: mean={v['mean']:+.3f} frac>0={v['frac_positive']:.2f}")

    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")
    args = ap.parse_args()
    main(tag="nmcal" if args.calibrate else "nm", calibrate=args.calibrate)
