"""Capstone: semantic-distance analysis (canonical version, ships with the capstone deposit).

Question (from paper 4's reviewer): is the persona-induced self-report shift
partly item-level instruction matching? IPIP agreeableness items nearly
paraphrase the persona text, so the say-effect could track semantic proximity
to the prompt rather than trait induction.

Method: embed each IPIP-50 item and each persona system text (all-minilm via
local Ollama); per item, compute the shift in mean scored answer from L0 (no
persona) to L4 (strongest clamp), pooled over the 18 configurations and both
contexts of the paper 4 collection. Analyses across the 50 items, per target:
  (1) r(similarity, |shift|), with Fisher 95% CI, plus on-/off-target splits;
  (2) OLS |shift| ~ similarity + on_trait (does wording predict shift once
      the measured construct is known?);
  (3) signed test: parroting predicts raw agreement rises most on items most
      similar to the prompt, so correlate similarity with the signed shift in
      RAW answers (no reverse-keying), r(similarity, raw L4 - raw L0).

Outputs: paper4_out/semantic_distance.csv (with signed_raw_shift column),
semantic_distance_summary.txt, fig_semantic_distance.png, persona_systems.txt
(the two embedded persona texts, for the capstone deposit).
"""
import json
import sqlite3
import sys
import urllib.request

sys.path.insert(0, "src")
sys.path.insert(0, ".")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yaml
from math import sqrt, log, exp

from llmpsy.instrument import load_instrument
from llmpsy.scoring import score_items

DB = "results/persona_full_merged.db"
EMBED_URL = "http://192.168.1.178:11434/api/embed"
EMBED_MODEL = "all-minilm"   # sentence-transformers all-MiniLM-L6-v2 via Ollama
TARGET_TRAIT = {"AGR_lo": "AGR", "CSN_hi": "CSN"}


def embed(texts):
    req = urllib.request.Request(
        EMBED_URL, data=json.dumps({"model": EMBED_MODEL, "input": texts}).encode(),
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)["embeddings"]


def cos(a, b):
    num = sum(x * y for x, y in zip(a, b))
    return num / (sqrt(sum(x * x for x in a)) * sqrt(sum(x * x for x in b)))


def pearson(xs, ys):
    n = len(xs)
    mx, my = sum(xs) / n, sum(ys) / n
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    dx = sqrt(sum((x - mx) ** 2 for x in xs))
    dy = sqrt(sum((y - my) ** 2 for y in ys))
    return num / (dx * dy) if dx and dy else float("nan")


def spearman(xs, ys):
    def rank(v):
        # average ranks for ties (canonical tie handling)
        order = sorted(range(len(v)), key=lambda i: v[i])
        r = [0.0] * len(v)
        i = 0
        while i < len(order):
            j = i
            while j + 1 < len(order) and v[order[j + 1]] == v[order[i]]:
                j += 1
            avg = (i + j) / 2.0
            for k in range(i, j + 1):
                r[order[k]] = avg
            i = j + 1
        return r
    return pearson(rank(xs), rank(ys))


def fisher_ci(r, n):
    z = 0.5 * log((1 + r) / (1 - r))
    se = 1 / sqrt(n - 3)
    f = lambda z: (exp(2 * z) - 1) / (exp(2 * z) + 1)
    return f(z - 1.96 * se), f(z + 1.96 * se)


def ols(y, X):
    """OLS with intercept; returns coefs, standard errors, t stats."""
    X = np.column_stack([np.ones(len(y))] + X)
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    resid = y - X @ beta
    mse = resid @ resid / (len(y) - X.shape[1])
    se = np.sqrt(np.diag(mse * np.linalg.inv(X.T @ X)))
    return beta, se, beta / se


inst = load_instrument("ipip50")
items = {it.id: it for it in inst.items}
personas = yaml.safe_load(open("data/personas.yaml"))["targets"]

ids = [it.id for it in inst.items]
vecs = embed([items[i].text for i in ids])
ivec = dict(zip(ids, vecs))
pvec = {t: embed([personas[t]["system"]])[0] for t in TARGET_TRAIT}

with open("paper4_out/persona_systems.txt", "w") as f:
    for t in TARGET_TRAIT:
        f.write(f"== {t} (embedded system text) ==\n{personas[t]['system']}\n\n")

con = sqlite3.connect(f"file:{DB}?mode=ro", uri=True)
df = pd.read_sql(
    "SELECT model, reasoning_mode, item_id, answer, clamp_level, persona_target "
    "FROM runs WHERE probe_type='selfreport' AND instrument='ipip50' "
    "AND model != 'olmo-3:7b' AND clamp_level IN ('L0','L4')", con)
con.close()
df = score_items(df)   # adds reverse-keyed 'scored'

base = df[(df.clamp_level == "L0")].groupby("item_id").scored.mean()
base_raw = df[(df.clamp_level == "L0")].groupby("item_id").answer.mean()

rows = []
for target, trait in TARGET_TRAIT.items():
    sel = df[(df.clamp_level == "L4") & (df.persona_target == target)]
    l4 = sel.groupby("item_id").scored.mean()
    l4_raw = sel.groupby("item_id").answer.mean()
    for iid in ids:
        rows.append(dict(
            target=target, item=iid, domain=items[iid].domain,
            on_trait=int(items[iid].domain == trait),
            similarity=round(cos(ivec[iid], pvec[target]), 4),
            abs_shift=round(abs(l4[iid] - base[iid]), 4),
            signed_raw_shift=round(l4_raw[iid] - base_raw[iid], 4)))
tab = pd.DataFrame(rows)
tab.to_csv("paper4_out/semantic_distance.csv", index=False)

lines = []
for target in TARGET_TRAIT:
    t = tab[tab.target == target]
    r_all = pearson(list(t.similarity), list(t.abs_shift))
    lo, hi = fisher_ci(r_all, len(t))
    rho_all = spearman(list(t.similarity), list(t.abs_shift))
    on = t[t.on_trait == 1]; off = t[t.on_trait == 0]
    r_on = pearson(list(on.similarity), list(on.abs_shift))
    r_off = pearson(list(off.similarity), list(off.abs_shift))
    lines.append(f"[{target}] all 50 items: r={r_all:+.3f} 95%CI[{lo:+.2f},{hi:+.2f}] "
                 f"(rho={rho_all:+.3f}); on-trait 10: r={r_on:+.3f}; off-trait 40: r={r_off:+.3f}")
    lines.append(f"[{target}] mean similarity on-trait={on.similarity.mean():.3f} "
                 f"off-trait={off.similarity.mean():.3f}; "
                 f"mean |shift| on-trait={on.abs_shift.mean():.3f} off-trait={off.abs_shift.mean():.3f} "
                 f"(ratio {on.abs_shift.mean()/off.abs_shift.mean():.2f}x)")
    beta, se, tstat = ols(t.abs_shift.to_numpy(),
                          [t.similarity.to_numpy(), t.on_trait.to_numpy(float)])
    lines.append(f"[{target}] OLS |shift| ~ sim + on_trait: "
                 f"sim beta={beta[1]:+.3f} (t={tstat[1]:+.2f}); "
                 f"on_trait beta={beta[2]:+.3f} (t={tstat[2]:+.2f})")
    r_signed = pearson(list(t.similarity), list(t.signed_raw_shift))
    slo, shi = fisher_ci(r_signed, len(t))
    r_signed_off = pearson(list(off.similarity), list(off.signed_raw_shift))
    lines.append(f"[{target}] SIGNED (raw agreement shift vs similarity): "
                 f"r={r_signed:+.3f} 95%CI[{slo:+.2f},{shi:+.2f}]; off-target only r={r_signed_off:+.3f}")
open("paper4_out/semantic_distance_summary.txt", "w").write("\n".join(lines) + "\n")
print("\n".join(lines))

fig, axes = plt.subplots(1, 2, figsize=(10, 4.2), sharey=True)
for ax, target in zip(axes, TARGET_TRAIT):
    t = tab[tab.target == target]
    for flag, color, label in [(1, "#D65F5F", "target-trait items"), (0, "#4878CF", "other items")]:
        s = t[t.on_trait == flag]
        ax.scatter(s.similarity, s.abs_shift, c=color, s=34, label=label, zorder=3)
    ax.set_title(f"persona: {target}")
    ax.set_xlabel("item-to-persona cosine similarity")
    ax.grid(alpha=0.25)
axes[0].set_ylabel("|item score shift| L0 to L4 (scale points)")
axes[0].legend(fontsize=8)
fig.suptitle("Does the induced shift track semantic similarity to the persona text?")
fig.tight_layout()
fig.savefig("paper4_out/fig_semantic_distance.png", dpi=200)
print("wrote paper4_out/semantic_distance.{csv,txt,png} + persona_systems.txt")