"""Results of the comparative-positioning run: evidence/records.jsonl -> evidence/results.json,
buyers.csv.

    python research/studies/market-perception/comparative-positioning/analyze.py [--evidence DIR]

Contrasts are within-buyer: every arm of a world branches from the same pre-exposure
conversation, so a buyer's arms differ only in what the exposure showed. Intervals are
buyer-clustered percentile bootstraps (``measure.paired_difference``) or Wilson
intervals for shares. Missing answers are dropped from the contrast they belong to.
"""

import argparse
import csv
import json
import pathlib
import sys

HERE = pathlib.Path(__file__).resolve().parent
ROOT = HERE.parents[3]
sys.path.insert(0, str(ROOT / "scripts/society"))

import factors  # noqa: E402
import measure  # noqa: E402

SPEC = json.loads((HERE / "spec.json").read_text())
OURS, THEIRS = SPEC["ours"], SPEC["theirs"]
DEFER = SPEC["options"][2]
REPS = 2000
CRIT = [f"weight_{c.replace(' ', '_')}" for c in SPEC["criteria"]]


def build(rows):
    """(seed, buyer, world, arm) -> values; pre-exposure values copied into every arm."""
    cells, buyers = {}, {}
    for r in rows:
        key = (r["seed"], r["buyer"], r["world"], r["arm"])
        c = cells.setdefault(key, {})
        c[f"{r['track']}.{r['stage']}.{r['item']}"] = r["value"]
        buyers[r["buyer"]] = {k: r[k] for k in ("seed", "need", "incumbent", "role")}
    for c in cells.values():
        for stage in ("pre", "post"):
            weights = [c.get(f"measure.{stage}.{k}") for k in CRIT]
            total = sum(w for w in weights if w is not None) if None not in weights else None
            for k, w in zip(CRIT, weights):
                c[f"{stage}.{k}"] = None if not total else 100.0 * w / total
            choice = c.get(f"choice.{stage}.choose")
            c[f"{stage}.ours"] = None if choice is None else choice == OURS
            c[f"{stage}.theirs"] = None if choice is None else choice == THEIRS
            c[f"{stage}.defer"] = None if choice is None else choice == DEFER
            c[f"{stage}.preference"] = c.get(f"measure.{stage}.relative_preference")
        pre, post = c.get("choice.pre.choose"), c.get("choice.post.choose")
        c["switch"] = None if pre is None or post is None else (
            "none" if pre == post else ("to_ours" if post == OURS else ("to_theirs" if post == THEIRS else "to_defer")))
        c["latency_shift"] = None if c.get("post.weight_query_latency") is None or c.get("pre.weight_query_latency") is None \
            else c["post.weight_query_latency"] - c["pre.weight_query_latency"]
        c["trial"] = c.get("choice.post.trial")
        c["source_independence"] = c.get("measure.post.source_independence")
    return cells, buyers


def num(v):
    return (1.0 if v else 0.0) if isinstance(v, bool) else v


def share(cells, metric, world, arm, need=None, buyers=None):
    vals = [num(c.get(metric)) for (s, b, w, a), c in cells.items() if w == world and a == arm
            and (need is None or buyers[b]["need"] == need)]
    vals = [v for v in vals if v is not None]
    out = {"mean": measure.mean(vals), "n": len(vals)}
    if vals and all(v in (0.0, 1.0) for v in vals):
        out["successes"] = int(sum(vals))
        out["lo"], out["hi"] = measure.wilson(int(sum(vals)), len(vals))
    return out


def contrast(cells, metric, world, arm, ref="none", need=None, buyers=None):
    pairs = []
    for (s, b, w, a), c in cells.items():
        if w == world and a == arm and (need is None or buyers[b]["need"] == need):
            pairs.append((num(c.get(metric)), num(cells.get((s, b, w, ref), {}).get(metric))))
    return measure.paired_difference(pairs, reps=REPS, seed=0, name=f"{metric}/{world}/{arm}/{ref}")


def interaction(cells, metric, arm, need=None, buyers=None):
    """(arm - none in baseline) - (arm - none in reversed), per buyer."""
    diffs = []
    for (s, b, w, a), c in cells.items():
        if w != "baseline" or a != arm or (need and buyers[b]["need"] != need):
            continue
        vals = [num(cells.get((s, b, ww, aa), {}).get(metric)) for ww, aa in
                (("baseline", arm), ("baseline", "none"), ("reversed", arm), ("reversed", "none"))]
        diffs.append(None if None in vals else (vals[0] - vals[1]) - (vals[2] - vals[3]))
    used = [d for d in diffs if d is not None]
    est, lo, hi = measure.bootstrap(used, lambda xs: measure.mean(xs), reps=REPS, seed=0, name=f"interaction/{metric}/{arm}")
    return {"estimate": est, "lo": lo, "hi": hi, "n": len(used)}


def conditional_share(cells, world, arm):
    """Share choosing ours among buyers who chose a product (did not defer), with a buyer
    bootstrap interval."""
    items = [(num(c.get("post.ours")), num(c.get("post.defer"))) for (s, b, w, a), c in cells.items()
             if w == world and a == arm and c.get("post.ours") is not None]

    def stat(xs):
        decided = [o for o, d in xs if not d]
        return sum(decided) / len(decided) if decided else None
    est, lo, hi = measure.bootstrap(items, stat, reps=REPS, seed=0, name=f"conditional/{world}/{arm}")
    return {"estimate": est, "lo": lo, "hi": hi, "decided": sum(1 for o, d in items if not d), "n": len(items)}


def empty_replies(rows, arm_key):
    """Per arm: prompts answered with nothing usable (every item of the stage invalid,
    e.g. an empty reply after the model spent its token budget reasoning) and prompts
    with some items invalid, out of all prompts. Reported because an uneven rate across
    arms would confound the contrasts."""
    stages = {}
    for r in rows:
        k = (arm_key(r), r["seed"], r["buyer"], r["track"], r["stage"], r.get("price"))
        stages.setdefault(k, []).append(r["valid"])
    out = {}
    for k, v in stages.items():
        e = out.setdefault(k[0], {"prompts": 0, "empty": 0, "partial": 0})
        e["prompts"] += 1
        e["empty"] += not any(v)
        e["partial"] += any(v) and not all(v)
    return dict(sorted(out.items()))


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--evidence", default=str(HERE / "evidence"))
    args = parser.parse_args(argv)
    ev = pathlib.Path(args.evidence)
    rows = [json.loads(line) for line in open(ev / "records.jsonl")]
    run = json.loads((ev / "run.json").read_text())
    cells, buyers = build(rows)
    worlds, arms, needs = list(SPEC["worlds"]), SPEC["arms"], SPEC["population"]["strata"]["need"]
    invalid = {}
    for r in rows:
        if not r["valid"]:
            k = f"{r['track']}/{r['stage']}/{r['item']}"
            invalid[k] = invalid.get(k, 0) + 1
    res = {"study": "comparative-positioning", "synthetic": True,
           "label": "Simulated buyers (an LLM role-playing personas). Not a measurement of real buyers.",
           "run": run, "n_buyers": len(buyers), "n_records": len(rows), "invalid_answers": invalid,
           "n_invalid": sum(invalid.values()), "empty_replies": empty_replies(rows, lambda r: f"{r['world']}/{r['arm']}"), "worlds": SPEC["worlds"],
           "n_buyers_by_need": {n: sum(1 for b in buyers.values() if b["need"] == n) for n in needs}}
    res["primary"] = {"outcome": "post.ours", "contrast": "campaign - none, baseline world",
                      **contrast(cells, "post.ours", "baseline", "campaign"),
                      "per_seed": {}}
    for seed in run["seeds"]:
        sub = {k: v for k, v in cells.items() if k[0] == seed}
        res["primary"]["per_seed"][str(seed)] = contrast(sub, "post.ours", "baseline", "campaign")
    res["falsification"] = {"reversed": contrast(cells, "post.ours", "reversed", "campaign"),
                            "interaction": interaction(cells, "post.ours", "campaign")}
    metrics = ["pre.ours", "post.ours", "post.theirs", "post.defer", "post.preference", "post.weight_query_latency",
               "pre.weight_query_latency", "latency_shift", "post.weight_connector_coverage", "post.weight_price",
               "source_independence"]
    res["arms"] = {m: {w: {a: share(cells, m, w, a) for a in arms} for w in worlds} for m in metrics}
    res["contrasts"] = {m: {w: {a: contrast(cells, m, w, a) for a in arms if a != "none"} for w in worlds}
                        for m in metrics if not m.startswith("pre.")}
    res["interactions"] = {m: {a: interaction(cells, m, a) for a in arms if a != "none"}
                           for m in ("post.ours", "post.theirs", "post.defer", "post.preference", "latency_shift")}
    res["conditional"] = {w: {a: conditional_share(cells, w, a) for a in arms} for w in worlds}
    res["switching"] = {}
    for w in worlds:
        res["switching"][w] = {}
        for a in arms:
            counts = {}
            for (s, b, ww, aa), c in cells.items():
                if ww == w and aa == a:
                    key = f"{c.get('choice.pre.choose')} -> {c.get('choice.post.choose')}"
                    counts[key] = counts.get(key, 0) + 1
            res["switching"][w][a] = dict(sorted(counts.items(), key=lambda kv: -kv[1]))
    res["trial"] = {w: {a: {} for a in arms} for w in worlds}
    for (s, b, w, a), c in cells.items():
        t = c.get("trial") or "missing"
        res["trial"][w][a][t] = res["trial"][w][a].get(t, 0) + 1
    res["needs"] = {n: {"primary": contrast(cells, "post.ours", "baseline", "campaign", need=n, buyers=buyers),
                        "reversed": contrast(cells, "post.ours", "reversed", "campaign", need=n, buyers=buyers),
                        "benchmark_baseline": contrast(cells, "post.ours", "baseline", "benchmark", need=n, buyers=buyers),
                        "benchmark_reversed": contrast(cells, "post.ours", "reversed", "benchmark", need=n, buyers=buyers),
                        "pre_ours": {w: share(cells, "pre.ours", w, "none", need=n, buyers=buyers) for w in worlds},
                        "post_ours": {w: {a: share(cells, "post.ours", w, a, need=n, buyers=buyers) for a in arms} for w in worlds}}
                    for n in needs}
    # Factor ablation: predict choosing ours after exposure, held out by (world, arm) scenario.
    abl = []
    for (s, b, w, a), c in cells.items():
        row = {"outcome": None if c.get("post.ours") is None else int(c["post.ours"]), "scenario": f"{w}/{a}",
               "buyer": f"{s}/{b}", "pre_ours": num(c.get("pre.ours")), "preference": c.get("post.preference"),
               "w_latency": c.get("post.weight_query_latency"), "w_connectors": c.get("post.weight_connector_coverage"),
               "w_price": c.get("post.weight_price"), "independence": c.get("source_independence"),
               "true_advantage_ours": int(w == "baseline"), "need_latency": int(buyers[b]["need"] == "latency"),
               "need_budget": int(buyers[b]["need"] == "budget")}
        row["latency_x_advantage"] = None if row["w_latency"] is None else row["w_latency"] * (1 if w == "baseline" else -1)
        abl.append(row)
    baseline = ["pre_ours", "preference"]
    expanded = baseline + ["w_latency", "w_connectors", "w_price", "independence", "true_advantage_ours",
                           "latency_x_advantage", "need_latency", "need_budget"]
    res["ablation"] = factors.ablation(abl, "outcome", baseline, expanded, "scenario", "buyer", reps=1000, seed=0)
    res["ablation"].update(baseline=baseline, expanded=expanded,
                           held_out="leave-one-scenario-out over the six (world, exposure) scenarios; predictors from the measurement track, outcome from the choice track")
    res["incremental"] = factors.incremental_value(abl, "outcome", baseline, [f for f in expanded if f not in baseline],
                                                   "scenario", "buyer", reps=300, seed=0)
    examples = []
    for (s, b, w, a), c in sorted(cells.items()):
        if w == "reversed" and a == "campaign" and c.get("choice.pre.choose") == OURS and len(examples) < 4:
            examples.append({"buyer": b, "need": buyers[b]["need"], "pre": c.get("choice.pre.choose"),
                             "pre_why": c.get("choice.pre.why"), "post": c.get("choice.post.choose"),
                             "post_why": c.get("choice.post.why"), "trial": c.get("trial")})
    for (s, b, w, a), c in sorted(cells.items()):
        if w == "baseline" and a == "campaign" and c.get("choice.post.choose") == DEFER and len(examples) < 7:
            examples.append({"buyer": b, "need": buyers[b]["need"], "world": w, "pre": c.get("choice.pre.choose"),
                             "pre_why": c.get("choice.pre.why"), "post": c.get("choice.post.choose"),
                             "post_why": c.get("choice.post.why"), "trial": c.get("trial")})
    res["examples"] = examples
    (ev / "results.json").write_text(json.dumps(res, indent=1, sort_keys=True) + "\n")
    cols = ["seed", "buyer", "need", "incumbent", "world", "arm", "choice.pre.choose", "choice.post.choose", "trial",
            "pre.preference", "post.preference", "pre.weight_query_latency", "post.weight_query_latency",
            "post.weight_connector_coverage", "post.weight_price", "source_independence", "switch"]
    with open(ev / "buyers.csv", "w", newline="") as f:
        wr = csv.writer(f)
        wr.writerow(cols)
        for (s, b, w, a), c in sorted(cells.items()):
            row = dict(c, seed=s, buyer=b, world=w, arm=a, need=buyers[b]["need"], incumbent=buyers[b]["incumbent"])
            wr.writerow(["" if row.get(k) is None else row.get(k) for k in cols])
    p = res["primary"]
    print(f"buyers {len(buyers)}, invalid {res['n_invalid']}")
    print(f"primary campaign-none baseline: {p['estimate']:+.3f} [{p['lo']:+.3f}, {p['hi']:+.3f}]")
    f = res["falsification"]
    print(f"reversed: {f['reversed']['estimate']:+.3f} [{f['reversed']['lo']:+.3f}, {f['reversed']['hi']:+.3f}]; "
          f"interaction {f['interaction']['estimate']:+.3f} [{f['interaction']['lo']:+.3f}, {f['interaction']['hi']:+.3f}]")


if __name__ == "__main__":
    main()