"""phx_composition.py — the impression-invariant test of the released model.

The audience-size critique of a raw-count backtest is correct: like counts conflate
"good post" with "big account". Public view counts would fix it, but the zero-API
syndication endpoint does not expose views (verified: payload has favorite_count and
conversation_count only). So we use a quantity that cancels impressions ALGEBRAICALLY
instead of statistically.

For a post shown N times:
    actual reply share  = replies / (replies + likes)          <- N cancels
    model  reply share  = P(reply) / (P(reply) + P(favorite))  <- per-impression by construction

Both sides are rates, not totals. If the model has ANY real per-impression skill, it
should predict which posts skew reply-heavy versus like-heavy, regardless of reach.
This is the strongest test of the null result that public data permits.

Usage: phxenv/Scripts/python phx_composition.py
"""
from __future__ import annotations

import json
import math

import numpy as np


def spearman(a, b):
    ra = np.argsort(np.argsort(a)).astype(np.float64)
    rb = np.argsort(np.argsort(b)).astype(np.float64)
    ra -= ra.mean(); rb -= rb.mean()
    d = math.sqrt((ra @ ra) * (rb @ rb))
    return float(ra @ rb / d) if d else 0.0


def bootstrap_ci(x, y, n=2000, seed=11):
    rng = np.random.default_rng(seed)
    stats = []
    idx = np.arange(len(x))
    for _ in range(n):
        s = rng.choice(idx, size=len(idx), replace=True)
        if np.ptp(x[s]) > 0 and np.ptp(y[s]) > 0:
            stats.append(spearman(x[s], y[s]))
    return (round(float(np.percentile(stats, 2.5)), 3),
            round(float(np.percentile(stats, 97.5)), 3))


def load(scores_npz, fetched):
    z = np.load(scores_npz)
    probs, pids = z["probs"], z["post_ids"].astype(str)
    auth = z["author_ids"].astype(str) if "author_ids" in z else None
    idx = {p: i for i, p in enumerate(pids)}
    rows = []
    with open(fetched, encoding="utf-8") as f:
        for line in f:
            try:
                r = json.loads(line)
            except Exception:
                continue
            if r.get("status") == "ok" and r["id"] in idx:
                rows.append((r, idx[r["id"]]))
    return probs, auth, rows


out = {"test": "impression-invariant engagement composition",
       "why": "view counts are unavailable from the zero-API path (syndication payload "
              "exposes favorite_count + conversation_count only, no views, no follower "
              "count). Composition cancels impressions algebraically instead."}

for label, (sn, fj) in {
    "random_sample": ("scores.npz", "fetched.jsonl"),
    "prolific_authors": ("scores_authors.npz", "fetched_authors.jsonl"),
}.items():
    probs, auth, rows = load(sn, fj)
    # engagement floor: composition is meaningless on near-zero engagement
    keep = [(r, i) for r, i in rows if (r["favs"] + r["replies"]) >= 10]
    if len(keep) < 50:
        out[label] = {"n": len(keep), "note": "too few posts above engagement floor"}
        continue
    act = np.array([r["replies"] / (r["replies"] + r["favs"]) for r, _ in keep])
    mod = np.array([probs[i, 4] / (probs[i, 4] + probs[i, 1]) for _, i in keep])
    rho = spearman(mod, act)
    lo, hi = bootstrap_ci(mod, act)
    res = {"n_posts": len(keep),
           "spearman_model_vs_actual_reply_share": round(rho, 3),
           "bootstrap_95ci": [lo, hi],
           "significant": not (lo <= 0 <= hi),
           "actual_reply_share_median": round(float(np.median(act)), 4),
           "model_reply_share_median": round(float(np.median(mod)), 5)}

    # within-author (audience held constant AND impressions cancelled)
    if auth is not None:
        by = {}
        for r, i in keep:
            by.setdefault(auth[i], []).append((r, i))
        cors = []
        for a, items in by.items():
            if len(items) < 6:
                continue
            aa = np.array([r["replies"] / (r["replies"] + r["favs"]) for r, _ in items])
            mm = np.array([probs[i, 4] / (probs[i, 4] + probs[i, 1]) for _, i in items])
            if np.ptp(aa) > 0 and np.ptp(mm) > 0:
                cors.append(spearman(mm, aa))
        if cors:
            c = np.array(cors)
            res["within_author"] = {"n_authors": len(c),
                                    "median_spearman": round(float(np.median(c)), 3),
                                    "frac_positive": round(float((c > 0).mean()), 3)}
    out[label] = res

# baseline: does a trivial heuristic beat the model on the same task?
probs, _, rows = load("scores.npz", "fetched.jsonl")
keep = [(r, i) for r, i in rows if (r["favs"] + r["replies"]) >= 10]
act = np.array([r["replies"] / (r["replies"] + r["favs"]) for r, _ in keep])
mod = np.array([probs[i, 4] / (probs[i, 4] + probs[i, 1]) for _, i in keep])
qm = np.array([1.0 if "?" in r["text"] else 0.0 for r, _ in keep])
out["baseline_comparison"] = {
    "note": "same task, same posts: can a one-line heuristic (does the text contain '?') "
            "predict reply share better than the released model?",
    "model_spearman": round(spearman(mod, act), 3),
    "question_mark_heuristic_spearman": round(spearman(qm, act), 3),
    "n": len(keep),
}

print(json.dumps(out, indent=1))
with open("composition_stats.json", "w", encoding="utf-8") as f:
    json.dump(out, f, indent=1)