"""Comparative positioning: two fictional historians, three exposures, two worlds.

    python research/studies/market-perception/comparative-positioning/run.py --phase final            # asks the model
    python research/studies/market-perception/comparative-positioning/run.py --phase final --offline  # replays

Each buyer (``spec.json`` population, stratified by need: latency, connectors, budget)
is run in both worlds: ``baseline`` (our product, Halvard, is faster on query latency,
the feature our positioning emphasises) and ``reversed`` (the competitor is faster;
nothing else changes). In each world the buyer reads both vendors' pages and chooses
(pre), then one of three exposures (``none``, an independent ``benchmark``, Halvard's
factual latency ``campaign``) is a branch of that same conversation, and the buyer
chooses again (post). The pre-exposure prompt is byte-identical across arms, so it is
asked once and every arm starts from the same pre choice (common random numbers).

Two tracks per buyer and world (``scripts/society/measure.py``, assumption 6):

choice track      pre: choose (choice)        post: choose, hands-on trial (choice), why
measurement track pre: criterion weights, relative preference (elicited)
                  post: criterion weights, relative preference, source independence (elicited)
"""

import argparse
import json
import pathlib
import sys
import time

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

import measure  # noqa: E402
import streams  # noqa: E402

SPEC = json.loads((HERE / "spec.json").read_text())
Item, Stage = measure.Item, measure.Stage
OURS, THEIRS = SPEC["ours"], SPEC["theirs"]


def pages(world, order_seed):
    latency = SPEC["worlds"][world]["latency"]
    order = streams.Stream(order_seed, "order/pages").shuffle([OURS, THEIRS])
    return "\n\n".join(SPEC["products"][p]["page"].format(latency=latency[p]) for p in order)


def criteria_items(when):
    return [Item.estimate(f"weight_{c.replace(' ', '_')}", f"points out of 100 you give to {c} in this decision "
                                                           "(all five weights should add up to 100)", 0, 100)
            for c in SPEC["criteria"]]


def preference_item():
    return Item.scale("relative_preference", f"which you lean toward: -3 = strongly {THEIRS}, 0 = no preference, "
                                             f"+3 = strongly {OURS}", lo=-3, hi=3)


def pre_choice(world, order_seed):
    return Stage("pre", "You are comparing two historians. Their vendors' pages:\n\n" + pages(world, order_seed)
                 + "\n\nIf you had to decide today, what would you choose?",
                 [Item.choice("choose", "your choice today", SPEC["options"]), Item.note("why", "why, in a few words")])


def pre_measure(world, order_seed):
    return Stage("pre", "You are comparing two historians. Their vendors' pages:\n\n" + pages(world, order_seed)
                 + "\n\nHow are you weighing this decision?", criteria_items("pre") + [preference_item()])


def exposure_text(world, arm):
    latency = SPEC["worlds"][world]["latency"]
    return SPEC["exposures"][arm].format(lat_ours=latency[OURS], lat_theirs=latency[THEIRS])


def post_choice(world, arm):
    return Stage("post", exposure_text(world, arm) + "\n\nWhat do you choose now?", [
        Item.choice("choose", "your choice now", SPEC["options"]),
        Item.choice("trial", "which product, if any, you would put through a hands-on trial before signing",
                    [OURS, THEIRS, "Both", "Neither"]),
        Item.note("why", "why, in a few words")])


def post_measure(world, arm):
    return Stage("post", exposure_text(world, arm) + "\n\nHow are you weighing this decision now?",
                 criteria_items("post") + [preference_item(),
                                           Item.scale("source_independence", "how independent of both vendors the "
                                                      "information you were most recently shown is (1 = not at all, "
                                                      "7 = fully independent)")])


def persona(buyer):
    return SPEC["persona"].format(**dict(buyer, need_text=SPEC["needs"][buyer["need"]]))


def buyers(seed, n):
    pop = SPEC["population"]
    return measure.sample_population(pop["sample"], n, seed, stream="buyers", strata=pop["strata"])


def run_one(journal, seed, buyer, world):
    order_seed = f"{seed}/{buyer['id']}"
    system = SPEC["system"] + "\n\nWho you are:\n" + persona(buyer)
    base = {"study": "comparative-positioning", "seed": seed, "world": world, "buyer": buyer["id"],
            "need": buyer["need"], "incumbent": buyer["incumbent"], "role": buyer["role"]}
    rows, transcripts = [], []
    for track, pre, post in (("choice", pre_choice, post_choice), ("measure", pre_measure, post_measure)):
        trunk = measure.Session(journal, system, order_seed)
        trunk.ask(pre(world, order_seed))
        for arm in SPEC["arms"]:
            branch = trunk.branch()
            branch.ask(post(world, arm))
            rows += measure.records(branch, dict(base, arm=arm, track=track))
            transcripts.append({"track": track, "arm": arm, "messages": branch.messages})
    return rows, dict(base, sessions=transcripts)


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--phase", choices=["pilot", "final"], default="final")
    parser.add_argument("--offline", action="store_true", help="answer only from the journal")
    parser.add_argument("--journal")
    parser.add_argument("--out")
    parser.add_argument("--threads", type=int, default=24)
    args = parser.parse_args(argv)
    out = pathlib.Path(args.out or (HERE / "evidence" if args.phase == "final"
                                    else ROOT / "target/market-perception/comparative-positioning-pilot"))
    out.mkdir(parents=True, exist_ok=True)
    journal_path = pathlib.Path(args.journal or out / "journal.jsonl")
    seeds, n = SPEC["seeds"][args.phase], SPEC["buyers_per_seed"][args.phase]
    import computerworld
    pin = {"study": "comparative-positioning", "spec": measure.digest(SPEC), "phase": args.phase}
    journal = measure.open_journal(journal_path, pin, offline=args.offline)
    jobs = [(seed, b, world) for seed in seeds for b in buyers(seed, n) for world in SPEC["worlds"]]
    prompts = len(jobs) * 2 * (1 + len(SPEC["arms"]))
    if prompts > SPEC["budget"]["model_calls_max"]:
        raise SystemExit(f"{prompts} prompts exceed the stated budget of {SPEC['budget']['model_calls_max']}")
    started = time.time()
    results = measure.run_all(jobs, lambda job: run_one(journal, *job), threads=args.threads)
    rows = [r for rs, _ in results for r in rs]
    with open(out / "records.jsonl", "w") as f:
        for r in rows:
            f.write(json.dumps(r, sort_keys=True) + "\n")
    with open(out / "transcripts.jsonl", "w") as f:
        for (_, t), (seed, b, world) in zip(results, jobs):
            if b["index"] < 3:
                f.write(json.dumps(t, sort_keys=True) + "\n")
    run = {"study": "comparative-positioning", "phase": args.phase, "seeds": seeds, "buyers_per_seed": n,
           "arms": SPEC["arms"], "worlds": list(SPEC["worlds"]), "jobs": len(jobs), "records": len(rows),
           "spec_digest": pin["spec"], "engine": computerworld.engine_version, "model": journal.model,
           "temperature": journal.client.temperature, "reasoning_effort": "low",
           "prompt_version": measure.PROMPT_VERSION, "offline": args.offline, "journal": journal.stats(),
           "seconds": round(time.time() - started, 1)}
    (out / ("run.json" if not args.offline else "replay.json")).write_text(json.dumps(run, indent=2) + "\n")
    print(json.dumps(run, indent=2))


if __name__ == "__main__":
    main()
