"""Brand priors: one product, six corporate identities, simulated industrial buyers.

    python research/studies/market-perception/brand-priors/run.py --phase final            # asks the model
    python research/studies/market-perception/brand-priors/run.py --phase final --offline  # replays the journal

Each buyer (``spec.json`` population, stratified by segment) meets the focal product
under every identity arm, in two separate sessions per arm so ratings cannot steer
choices (``scripts/society/measure.py``, assumption 6):

choice track      1 discovery   which of four directory listings to open (choice)
                  2 company     a forwarded company page + product sheet; next step (choice)
                  3 price x3    side-by-side comparison at $1,190 / $1,490 / $1,890 (choice),
                                three branches of one conversation
                  4 experience  an identical six-month report, after the $1,490 branch:
                                satisfaction (elicited) and reorder (choice)
measurement track 1 name only   institutional and premium associations (elicited)
                  2 page        competence, credibility, risk, viability ... (elicited)

Everything but the identity block (name, tagline, positioning copy) is byte-identical
across arms, and listing and option order depend on the buyer alone. Budget: 8 prompts
per buyer per arm. Writes evidence/records.jsonl, evidence/transcripts.jsonl and
evidence/run.json; analyze.py turns those into results.
"""

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
NONE = "None of these yet"


def identity(arm):
    ident = SPEC["identities"][arm]
    copy = SPEC["copy"][ident["copy"]].format(name=ident["name"])
    return ident["name"], ident["tagline"], copy


def company_block(arm):
    name, tagline, copy = identity(arm)
    return f"{name}\n{tagline}\n\n{copy}\n{SPEC['company_facts']}\n\nProduct:\n{SPEC['product']['sheet']}"


def focal_label(arm):
    return f"{SPEC['product']['model']} from {identity(arm)[0]}"


def listing_order(order_seed):
    return streams.Stream(order_seed, "order/listings").shuffle(["focal", "tarnfield", "bexa", "distractor"])


def product_order(order_seed):
    return streams.Stream(order_seed, "order/products").shuffle(["focal", "tarnfield", "bexa"])


def discovery_stage(arm, order_seed):
    name, tagline, _ = identity(arm)
    entries = {"focal": (name, tagline), "distractor": (SPEC["distractor"]["name"], SPEC["distractor"]["tagline"]),
               **{k: (v["name"], v["tagline"]) for k, v in SPEC["competitor_listings"].items()}}
    order = listing_order(order_seed)
    lines = [f"- {entries[slot][0]}: {entries[slot][1]}" for slot in order]
    stimulus = ("You search an industrial supplier directory for \"industrial edge gateway\". It returns:\n"
                + "\n".join(lines) + "\n\nWhich entries, if any, do you open to read more?")
    return Stage("discovery", stimulus, [Item.choice("open", "the entries you open", [entries[s][0] for s in order],
                                                     many=True, shuffle=False)])


def company_stage(arm):
    stimulus = ("Whether or not you opened it, a colleague forwards you this supplier's page:\n\n"
                + company_block(arm) + "\n\nWhat do you do next?")
    return Stage("company", stimulus, [
        Item.choice("next", "your next step", ["Request a quote", "Save the datasheet for later", "Move on"], shuffle=False),
        Item.note("why", "why, in a few words")])


def price_stage(arm, price, order_seed):
    rows = {"focal": f"{focal_label(arm)} | ${price:,} per unit | as on the page you read",
            **{k: f"{v['label']} | ${v['price']:,} per unit | {v['sheet']}" for k, v in SPEC["competitors"].items()}}
    labels = {"focal": focal_label(arm), **{k: v["label"] for k, v in SPEC["competitors"].items()}}
    order = product_order(order_seed)
    stimulus = ("You now compare the options side by side for four units:\n"
                + "\n".join("- " + rows[s] for s in order)
                + "\n\nWhich do you choose to buy now?")
    return Stage(f"price-{price}", stimulus, [
        Item.choice("choose", "what you buy", [labels[s] for s in order] + [NONE], shuffle=False),
        Item.note("why", "why, in a few words")])


def experience_stage(arm):
    name = identity(arm)[0]
    return Stage("experience", SPEC["experience"].format(name=name), [
        Item.scale("satisfaction", "how satisfied you are with the EG-400 overall (1 = very dissatisfied, 7 = very satisfied)"),
        Item.decision("reorder", f"whether you would order the EG-400 from {name} again for your next site")])


def name_stage(arm):
    name = identity(arm)[0]
    return Stage("name", f"Before anything else: you see the company name \"{name}\" for the first time, with nothing else about it.", [
        Item.scale("name_institutional_association", "how strongly the name alone suggests a research institute, laboratory or university spin-out (1-7)"),
        Item.scale("name_premium_association", "how strongly the name alone suggests a premium, design-led supplier (1-7)")])


MEASURE_ITEMS = [
    Item.scale("perceived_competence", "how capable this company is at building what it sells (1-7)"),
    Item.scale("institutional_credibility", "how far you trust it as an established, accountable institution (1-7)"),
    Item.scale("perceived_vendor_risk", "how risky it would be to depend on this supplier in production (1 = not risky, 7 = very risky)"),
    Item.estimate("expected_vendor_viability", "the percent chance it still sells and supports this product in five years", 0, 100),
    Item.scale("perceived_technical_depth", "how much real engineering depth is behind the product (1-7)"),
    Item.scale("expected_product_maturity", "how field-proven and free of early defects you expect it to be (1-7)"),
    Item.scale("expected_quality", "the overall quality you expect in your own deployment (1-7)"),
    Item.estimate("expected_failure_rate", "the percent of units you expect to need repair or replacement within three years", 0, 100),
    Item.scale("anticipated_maintenance_burden", "how much ongoing engineering time you expect it to take (1 = very little, 7 = a lot)"),
    Item.scale("procurement_friction", "how hard it would be to get this supplier approved and purchased (1 = easy, 7 = very hard)"),
    Item.scale("confidence_in_leadership", "confidence in the judgment of the people running the company (1-7)"),
    Item.estimate("willingness_to_pay", "the most you would pay per unit, in US dollars, before choosing an alternative", 0, 10000),
    Item.scale("overall_impression", "your overall impression of the company and product (1-7)"),
]


def page_stage(arm):
    return Stage("page", "You read this supplier's page:\n\n" + company_block(arm) + "\n\nRate it.", MEASURE_ITEMS)


def persona(buyer):
    values = dict(buyer, protocols=" and ".join(buyer["protocols"]))
    return SPEC["personas"][buyer["segment"]].format(**values)


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, arm):
    order_seed = f"{seed}/{buyer['id']}"
    system = SPEC["system"] + "\n\nWho you are:\n" + persona(buyer)
    base = {"study": "brand-priors", "seed": seed, "arm": arm, "buyer": buyer["id"], "segment": buyer["segment"],
            "risk": buyer["risk"], "budget": buyer["budget"], "industry": buyer["industry"],
            "protocols": buyer["protocols"]}
    rows, transcripts = [], []
    choice = measure.Session(journal, system, order_seed)
    choice.ask(discovery_stage(arm, order_seed))
    choice.ask(company_stage(arm))
    rows += measure.records(choice, dict(base, track="choice", price=None))
    trunk = len(choice.log)
    for price in SPEC["prices"]:
        branch = choice.branch()
        branch.ask(price_stage(arm, price, order_seed))
        if price == SPEC["base_price"]:
            branch.ask(experience_stage(arm))
        branch.log = branch.log[trunk:]
        rows += measure.records(branch, dict(base, track="choice", price=price))
        transcripts.append({"track": "choice", "price": price, "messages": branch.messages})
    rating = measure.Session(journal, system, order_seed)
    rating.ask(name_stage(arm))
    rating.ask(page_stage(arm))
    rows += measure.records(rating, dict(base, track="measure", price=None))
    transcripts.append({"track": "measure", "messages": rating.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", help="default: evidence/journal.jsonl (final) or target/ (pilot)")
    parser.add_argument("--out", help="default: evidence/ (final) or target/market-perception/brand-priors-pilot")
    parser.add_argument("--threads", type=int, default=24)
    parser.add_argument("--seeds", type=int, nargs="*", help="override the phase's seeds")
    args = parser.parse_args(argv)
    out = pathlib.Path(args.out or (HERE / "evidence" if args.phase == "final"
                                    else ROOT / "target/market-perception/brand-priors-pilot"))
    out.mkdir(parents=True, exist_ok=True)
    journal_path = pathlib.Path(args.journal or out / "journal.jsonl")
    seeds = args.seeds or SPEC["seeds"][args.phase]
    n = SPEC["buyers_per_seed"][args.phase]
    import runner  # noqa: F401  (the journal; fails clearly without the engine binding)
    import computerworld
    pin = {"study": "brand-priors", "spec": measure.digest(SPEC), "phase": args.phase}
    journal = measure.open_journal(journal_path, pin, offline=args.offline)
    jobs = [(seed, b, arm) for seed in seeds for b in buyers(seed, n) for arm in SPEC["arms"]]
    budget = SPEC["budget"]["model_calls_max"]
    if len(jobs) * 8 > budget:
        raise SystemExit(f"{len(jobs) * 8} prompts exceed the stated budget of {budget}")
    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, arm) in zip(results, jobs):
            if b["index"] < 3:
                f.write(json.dumps(t, sort_keys=True) + "\n")
    run = {"study": "brand-priors", "phase": args.phase, "seeds": seeds, "buyers_per_seed": n, "arms": SPEC["arms"],
           "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()
