"""Results of the brand-priors run: evidence/records.jsonl -> evidence/results.json, results.csv.

    python research/studies/market-perception/brand-priors/analyze.py [--evidence DIR]

Every number the page shows comes from results.json. Contrasts are within-buyer (each
buyer meets every arm) with buyer-clustered percentile bootstrap intervals
(``measure.paired_difference``); proportions also get Wilson intervals. Missing answers
drop the pair from that contrast and are counted, never imputed.
"""

import argparse
import csv
import json
import pathlib
import re
import statistics
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())
REPS = 2000
LATENTS = ["perceived_competence", "institutional_credibility", "perceived_vendor_risk", "expected_vendor_viability",
           "perceived_technical_depth", "expected_product_maturity", "expected_quality", "expected_failure_rate",
           "anticipated_maintenance_burden", "procurement_friction", "confidence_in_leadership", "willingness_to_pay",
           "overall_impression", "name_institutional_association", "name_premium_association"]
CATEGORY = {"gateway", "gateways", "industrial", "network", "networks", "systems", "products", "controls",
            "devices", "plant", "data", "automation"}


def name_features(name):
    """The derived name factors (factor registry: name_syllables, name_pronounceability,
    name_descriptiveness): deterministic functions of the letters."""
    words = re.findall(r"[a-z]+", name.lower().replace("inc", ""))
    syllables = 0
    for w in words:
        groups = re.findall(r"[aeiouy]+", w[:-1] if w.endswith("e") and len(w) > 2 else w)
        syllables += max(1, len(groups))
    letters = "".join(words)
    vowels = set("aeiouy")
    pairs = [(a in vowels) != (b in vowels) for a, b in zip(letters, letters[1:])]
    return {"name_syllables": syllables,
            "name_pronounceability": round(sum(pairs) / len(pairs), 3) if pairs else 0.0,
            "name_descriptiveness": any(w in CATEGORY for w in words)}


def focal_of(value, arm):
    if value is None:
        return None
    return value == f"{SPEC['product']['model']} from {SPEC['identities'][arm]['name']}"


def slot_of(value):
    if value is None:
        return None
    if value.startswith(SPEC["product"]["model"] + " from"):
        return "focal"
    for k, v in SPEC["competitors"].items():
        if value == v["label"]:
            return k
    return "none"


def table(rows):
    """(seed, buyer, arm) -> {item or item@price: value} and buyer -> attributes."""
    cells, buyers = {}, {}
    for r in rows:
        key = (r["seed"], r["buyer"], r["arm"])
        cell = cells.setdefault(key, {})
        name = r["item"] if r["stage"] not in ("experience",) and not r["stage"].startswith("price-") else \
            (f"{r['item']}@{r['price']}" if r["stage"].startswith("price-") else r["item"])
        if r["stage"] == "company" and r["item"] == "why":
            name = "why_company"
        cell[name] = r["value"]
        buyers[r["buyer"]] = {k: r[k] for k in ("seed", "segment", "risk", "budget", "industry", "protocols")}
    for (seed, buyer, arm), cell in cells.items():
        name = SPEC["identities"][arm]["name"]
        opened = cell.get("open")
        cell["listing_discovery"] = None if opened is None else name in opened
        cell["engagement_request"] = None if cell.get("next") is None else cell["next"] == "Request a quote"
        for price in SPEC["prices"]:
            cell[f"focal_choice@{price}"] = focal_of(cell.get(f"choose@{price}"), arm)
            cell[f"slot@{price}"] = slot_of(cell.get(f"choose@{price}"))
        cell["repurchase_choice"] = cell.get("reorder")
        s, q = cell.get("satisfaction"), cell.get("expected_quality")
        cell["expectation_gap"] = None if s is None or q is None else s - q
    return cells, buyers


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


def per_arm(cells, metric, arms, segment=None, buyers=None):
    out = {}
    for arm in arms:
        vals = [num(c.get(metric)) for (s, b, a), c in cells.items() if a == arm
                and (segment is None or buyers[b]["segment"] == segment)]
        vals = [v for v in vals if v is not None]
        entry = {"mean": measure.mean(vals), "n": len(vals)}
        if vals and all(v in (0.0, 1.0) for v in vals):
            entry["successes"] = int(sum(vals))
            entry["lo"], entry["hi"] = measure.wilson(int(sum(vals)), len(vals))
        elif len(vals) > 1:
            sd = statistics.stdev(vals)
            entry["sd"] = sd
            entry["lo"], entry["hi"] = (entry["mean"] - 1.96 * sd / len(vals) ** 0.5,
                                        entry["mean"] + 1.96 * sd / len(vals) ** 0.5)
        out[arm] = entry
    return out


def contrast(cells, metric, treated, control, segment=None, buyers=None, seed=None, name=None):
    pairs = []
    for (s, b, a), c in cells.items():
        if a != treated or (segment and buyers[b]["segment"] != segment) or (seed is not None and s != seed):
            continue
        other = cells.get((s, b, control), {})
        pairs.append((num(c.get(metric)), num(other.get(metric))))
    return measure.paired_difference(pairs, reps=REPS, seed=0, name=name or f"{metric}/{treated}/{control}")


def ablation_rows(cells, buyers):
    """One row per (buyer, arm, price) with the outcome and candidate predictors. The
    elicited latents come from the measurement track of the same buyer and arm (a separate
    session), the outcome from the choice track."""
    rows = []
    comp = sorted(v["price"] for v in SPEC["competitors"].values())
    for (s, b, arm), c in cells.items():
        for price in SPEC["prices"]:
            focal = c.get(f"focal_choice@{price}")
            ref = statistics.median(comp)
            row = {"outcome": None if focal is None else int(focal), "arm": arm, "buyer": b, "seed": s,
                   "price_ratio": price / ref,
                   "seg_budget": int(buyers[b]["segment"] == "budget"),
                   "seg_procurement": int(buyers[b]["segment"] == "procurement"),
                   "risk_low": int(buyers[b]["risk"] == "low"), "risk_high": int(buyers[b]["risk"] == "high"),
                   "budget_headroom": buyers[b]["budget"] / price}
            for latent in LATENTS:
                row[latent] = num(c.get(latent))
            rows.append(row)
    return rows


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())
    arms, control = SPEC["arms"], SPEC["control"]
    cells, buyers = table(rows)
    segments = SPEC["population"]["strata"]["segment"]
    invalid = {}
    for r in rows:
        if not r["valid"]:
            invalid[f"{r['stage']}/{r['item']}"] = invalid.get(f"{r['stage']}/{r['item']}", 0) + 1
    res = {"study": "brand-priors", "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: r["arm"]),
           "identities": {a: dict(SPEC["identities"][a], **name_features(SPEC["identities"][a]["name"])) for a in arms},
           "prices": SPEC["prices"], "base_price": SPEC["base_price"]}
    base = SPEC["base_price"]
    primary = f"focal_choice@{base}"
    res["primary"] = {"outcome": primary, "contrast": "kuvra - neutral",
                      **contrast(cells, primary, "kuvra", control),
                      "per_seed": {str(s): contrast(cells, primary, "kuvra", control, seed=s) for s in run["seeds"]}}
    res["negative_control"] = {"contrast": "neutral-placebo - neutral", **contrast(cells, primary, "neutral-placebo", control)}
    behaviour = ["listing_discovery", "engagement_request"] + [f"focal_choice@{p}" for p in SPEC["prices"]] + \
                ["repurchase_choice", "satisfaction", "expectation_gap"]
    res["arms"] = {m: per_arm(cells, m, arms) for m in behaviour + LATENTS}
    res["contrasts"] = {m: {a: contrast(cells, m, a, control) for a in arms if a != control} for m in behaviour + LATENTS}
    res["decomposition"] = {}
    for brand in ("kuvra", "orume"):
        res["decomposition"][brand] = {
            "name_alone": contrast(cells, primary, f"{brand}-name", control),
            "copy_given_name": contrast(cells, primary, brand, f"{brand}-name"),
            "total": contrast(cells, primary, brand, control)}
    res["choice_shares"] = {}
    for p in SPEC["prices"]:
        res["choice_shares"][str(p)] = {}
        for arm in arms:
            counts = {}
            for (s, b, a), c in cells.items():
                if a == arm:
                    counts[c.get(f"slot@{p}") or "missing"] = counts.get(c.get(f"slot@{p}") or "missing", 0) + 1
            res["choice_shares"][str(p)][arm] = counts
    res["segments"] = {seg: {"arms": per_arm(cells, primary, arms, seg, buyers),
                             "contrasts": {a: contrast(cells, primary, a, control, seg, buyers) for a in arms if a != control},
                             "competence": {a: contrast(cells, "perceived_competence", a, control, seg, buyers) for a in arms if a != control},
                             "risk": {a: contrast(cells, "perceived_vendor_risk", a, control, seg, buyers) for a in arms if a != control},
                             "n_buyers": sum(1 for b in buyers.values() if b["segment"] == seg)}
                       for seg in segments}
    # Presentation order: the focal product's position in the comparison is drawn per buyer
    # (the same in every arm); does position move choice?
    import importlib.util
    spec_run = importlib.util.spec_from_file_location("bp_run", HERE / "run.py")
    bp = importlib.util.module_from_spec(spec_run)
    spec_run.loader.exec_module(bp)
    position = {}
    for (s, b, a), c in cells.items():
        pos = bp.product_order(f"{s}/{b}").index("focal") + 1
        v = c.get(primary)
        if v is not None:
            position.setdefault(pos, []).append(int(v))
    res["position"] = {str(k): {"mean": measure.mean(v), "n": len(v), "lo": measure.wilson(sum(v), len(v))[0],
                                "hi": measure.wilson(sum(v), len(v))[1]} for k, v in sorted(position.items())}
    listing_pos = {}
    for (s, b, a), c in cells.items():
        pos = bp.listing_order(f"{s}/{b}").index("focal") + 1
        v = c.get("listing_discovery")
        if v is not None:
            listing_pos.setdefault(pos, []).append(int(v))
    res["listing_position"] = {str(k): {"mean": measure.mean(v), "n": len(v)} for k, v in sorted(listing_pos.items())}
    # Factor ablation: held out by identity arm.
    abl_rows = ablation_rows(cells, buyers)
    baseline = ["price_ratio", "overall_impression", "seg_budget", "seg_procurement"]
    expanded = baseline + [f for f in LATENTS if f != "overall_impression"] + ["risk_low", "risk_high", "budget_headroom"]
    res["ablation"] = factors.ablation(abl_rows, "outcome", baseline, expanded, "arm", "buyer", reps=1000, seed=0)
    res["ablation"]["baseline"], res["ablation"]["expanded"] = baseline, expanded
    res["ablation"]["held_out"] = "leave-one-identity-arm-out; rows are (buyer, arm, price); predictors from the measurement track, outcome from the choice track"
    candidates = [f for f in expanded if f not in baseline]
    res["incremental"] = factors.incremental_value(abl_rows, "outcome", baseline, candidates, "arm", "buyer", reps=300, seed=0)
    # Examples: the same buyer under control and Kuvra, with their stated reasons.
    examples = []
    for (s, b, a), c in sorted(cells.items()):
        if a != "kuvra" or len(examples) >= 6:
            continue
        n = cells.get((s, b, control), {})
        if c.get(primary) != n.get(primary) and c.get(primary) is not None and n.get(primary) is not None:
            examples.append({"buyer": b, "segment": buyers[b]["segment"], "risk": buyers[b]["risk"],
                             "budget": buyers[b]["budget"],
                             "neutral": {"choose": n.get(f"choose@{base}"), "why": n.get(f"why@{base}"),
                                         "competence": n.get("perceived_competence"), "risk": n.get("perceived_vendor_risk")},
                             "kuvra": {"choose": c.get(f"choose@{base}"), "why": c.get(f"why@{base}"),
                                       "competence": c.get("perceived_competence"), "risk": c.get("perceived_vendor_risk")}})
    res["examples"] = examples
    # Stated reasons at $1,490: how often a rationale names the supplier as young or unproven,
    # or names the competitor as established (simple keyword counts over the checked-in text).
    words = {"young_or_unproven": ("startup", "start-up", "unproven", "young", "new company", "2019"),
             "established": ("established", "proven", "track record", "longer history")}
    res["reasons"] = {}
    for arm in arms:
        texts = [str(c.get(f"why@{base}") or "").lower() for (s, b, a), c in cells.items() if a == arm]
        res["reasons"][arm] = {k: sum(any(w in t for w in ws) for t in texts) for k, ws in words.items()}
        res["reasons"][arm]["n"] = len(texts)
    res["reason_words"] = {k: list(v) for k, v in words.items()}
    res["n_buyers_by_segment"] = {seg: sum(1 for b in buyers.values() if b["segment"] == seg) for seg in segments}
    (ev / "results.json").write_text(json.dumps(res, indent=1, sort_keys=True, default=float) + "\n")
    with open(ev / "results.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["metric", "arm", "mean", "lo", "hi", "n", "diff_vs_neutral", "diff_lo", "diff_hi", "pairs"])
        for m, by in res["arms"].items():
            for arm, e in by.items():
                c = res["contrasts"][m].get(arm, {})
                w.writerow([m, arm, e["mean"], e.get("lo"), e.get("hi"), e["n"], c.get("estimate"), c.get("lo"), c.get("hi"), c.get("n")])
    columns = ["seed", "buyer", "arm", "segment", "risk", "budget", "industry", "listing_discovery", "next",
               "engagement_request"] + [f"{k}@{p}" for p in SPEC["prices"] for k in ("choose", "focal_choice")] + \
              ["satisfaction", "reorder", "expectation_gap"] + LATENTS
    with open(ev / "buyers.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(columns)
        for (s, b, a), c in sorted(cells.items()):
            row = dict(c, seed=s, buyer=b, arm=a, **{k: buyers[b][k] for k in ("segment", "risk", "budget", "industry")})
            w.writerow(["" if row.get(k) is None else row.get(k) for k in columns])
    p = res["primary"]
    print(f"buyers {len(buyers)}, records {len(rows)}, invalid {res['n_invalid']}")
    print(f"primary kuvra-neutral @${base}: {p['estimate']:+.3f} [{p['lo']:+.3f}, {p['hi']:+.3f}] n={p['n']}")
    nc = res["negative_control"]
    print(f"placebo: {nc['estimate']:+.3f} [{nc['lo']:+.3f}, {nc['hi']:+.3f}]")
    print("ablation log-loss reduction", res["ablation"]["log_loss_reduction"])


if __name__ == "__main__":
    main()