#!/usr/bin/env python3
"""Calibrate the ambiguous boundary: find a prompt whose filler space is
most evenly split between the two senses (max detection power), n=50 each."""

import json
from pathlib import Path

import numpy as np
from run_pilot import FIN_ANCHOR, RIV_ANCHOR, load_env_key, normalize
from run_pilot2 import first_sentence

OUT = Path(__file__).parent
CANDIDATES = {
    "C1": "He stood up, grabbed his coat, and walked down to the bank.",
    "C2": "After lunch, they walked along the path that led down to the bank.",
    "C3": "The children ran ahead of us, all the way down to the bank.",
    "C4": "Early that morning, she made her way down to the bank.",
    "C5": "We met them near the bank, just before sunset.",
    "C6": "He spent most of the afternoon down at the bank.",
}

from openai import OpenAI
client = OpenAI(api_key=load_env_key())

def embed(texts):
    resp = client.embeddings.create(model="text-embedding-3-small",
                                    input=[t if t.strip() else " " for t in texts])
    return np.array([d.embedding for d in resp.data])

fin_a, riv_a = normalize(embed([FIN_ANCHOR, RIV_ANCHOR]))
MID = -0.100  # midpoint of pilot-2 fs FIN/RIV reference means

report = {}
for key, prompt in CANDIDATES.items():
    resp = client.completions.create(model="gpt-3.5-turbo-instruct",
                                     prompt=prompt, max_tokens=40,
                                     temperature=1.0, n=50)
    conts = [c.text for c in resp.choices]
    with open(OUT / f"cal_completions_{key}.jsonl", "w") as f:
        for t in conts:
            f.write(json.dumps({"prompt": prompt, "completion": t}) + "\n")
    X = normalize(embed([first_sentence(t) for t in conts]))
    margins = X @ fin_a - X @ riv_a
    fin_n = int((margins > MID).sum())
    report[key] = {"prompt": prompt, "fin": fin_n, "riv": 50 - fin_n,
                   "margins": [float(m) for m in margins]}
    print(f"{key}: {fin_n:2d} fin / {50-fin_n:2d} riv   {prompt}", flush=True)

json.dump(report, open(OUT / "calibration.json", "w"), indent=2)
