#!/usr/bin/env python3
"""Every count in the paper, the consistency check, and the four SVG figures.

Inputs : data/corrections.csv, data/laws.csv, data/reversals.csv,
         data/messages_index.csv, data/codebook.json
Outputs: data/stats.json, data/tables.md, data/consistency_check.csv,
         <figdir>/fig1..fig4 .svg (only with --figdir; needs matplotlib)
"""
import argparse, collections, csv, json, math, os, random, re

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "..", "data")

def load():
    C = list(csv.DictReader(open(os.path.join(DATA, "corrections.csv"))))
    L = list(csv.DictReader(open(os.path.join(DATA, "laws.csv"))))
    V = list(csv.DictReader(open(os.path.join(DATA, "reversals.csv"))))
    M = {r["msg_id"]: r for r in csv.DictReader(open(os.path.join(DATA, "messages_index.csv")))}
    B = json.load(open(os.path.join(DATA, "codebook.json")))
    return C, L, V, M, B

def covering(laws, project, tell, date, strict=True):
    out = []
    for l in laws:
        if l["scope"] not in (project, "*"):
            continue
        if tell not in l["tells"].split(";"):
            continue
        if (l["date"] < date) if strict else (l["date"] <= date):
            out.append(l)
    return out

def law_status(laws, r):
    """in_force: a rule covering this tell applied to this project before this date.
    elsewhere: no such rule here, but one existed in another project.
    none: no written rule for this tell anywhere yet."""
    if covering(laws, r["project"], r["tell"], r["date"]):
        return "in_force"
    for l in laws:
        if l["scope"] != r["project"] and r["tell"] in l["tells"].split(";") and l["date"] < r["date"]:
            return "elsewhere"
    return "none"

KW = [  # independent keyword classifier for the consistency check (pattern text only)
    ("DOT", r"\bdot|pill|badge|indicator|marker line"),
    ("TYP", r"font|type|mono|serif|digits|figures|heading|label|eyebrow|wordmark in"),
    ("GLO", r"glow|glass|blur|sheen|shine|frost|fade under|shade|texture|reflection|neon"),
    ("COL", r"colou?r|purple|blue|white|palette|accent|navy|theme|cream|paper|green|dark mode"),
    ("ICO", r"icon|emoji|arrow"),
    ("IDN", r"logo|mark|wordmark|mascot"),
    ("MOT", r"animation|entrance|fade|motion|floating|loop|flicker|splash|cross-fade"),
    ("CLU", r"text|clutter|over-explain|copy|options|label|overwhelm|stressful|em dash"),
    ("NAV", r"sidebar|tab and|two .*bars|plus button|skip link|side navigation"),
    ("CMP", r"stock|default|keypad|numpad|spinner|snackbar|dropdown|native|onboarding|underline"),
    ("GIM", r"catalogue|customi[sz]ation|alternate|frames|demo|grid behind|options$|mosaic"),
    ("SRF", r"outline|edge|box|card effect|surface|shape|hairline|ring|depth|header"),
    ("TPL", r"section|card|dashboard|layout|space|hero|slideshow|tiles|presentation"),
    ("DED", r"dead|dull|basic|plain|flat|empty|alive|generic-looking|meh"),
    ("GEN", r"ai|generic|slop"),
]

def kw_code(text):
    t = text.lower()
    for code, pat in KW:
        if re.search(pat, t):
            return code
    return "GEN"

def kappa(a, b):
    n = len(a); po = sum(x == y for x, y in zip(a, b)) / n
    ca, cb = collections.Counter(a), collections.Counter(b)
    pe = sum(ca[k] * cb.get(k, 0) for k in ca) / (n * n)
    return po, (po - pe) / (1 - pe) if pe < 1 else 1.0

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--figdir"); a = ap.parse_args()
    C, L, V, M, B = load()
    tells = B["tells"]; EX, BL = set(B["excess_pole"]), set(B["bland_pole"])
    S = {}
    S["n_corrections"] = len(C)
    S["projects"] = sorted({r["project"] for r in C}); S["n_projects"] = len(S["projects"])
    S["date_range"] = [min(r["date"] for r in C), max(r["date"] for r in C)]
    S["n_messages_design"] = len(M)
    ev_msgs = {e for r in C for e in r["evidence"].split(";") if re.fullmatch(r"T\d{4}", e)}
    S["n_messages_coded"] = len(ev_msgs)
    S["by_tell"] = collections.Counter(r["tell"] for r in C).most_common()
    S["by_project"] = collections.Counter(r["project"] for r in C).most_common()
    S["projects_per_tell"] = {t: len({r["project"] for r in C if r["tell"] == t}) for t in tells}
    S["pole"] = {"excess": sum(r["tell"] in EX for r in C), "bland": sum(r["tell"] in BL for r in C),
                 "neither": sum(r["tell"] not in EX | BL for r in C)}
    src = collections.Counter()
    for r in C:
        for s in r["source"].split("+"): src[s] += 1
    S["by_source"] = dict(src)
    S["transcript_any"] = sum("transcript" in r["source"] for r in C)
    S["memory_only"] = sum(r["source"] == "memory" for r in C)
    S["multi_source"] = sum("+" in r["source"] for r in C)
    # strategies: coded + screenshot from attachments on evidence messages
    st = collections.Counter(); shot_auto = 0
    for r in C:
        ss = set(x for x in r["strategies"].split(";") if x)
        if any(M.get(e, {}).get("has_image") == "1" for e in r["evidence"].split(";")):
            if "SHOT" not in ss: shot_auto += 1
            ss.add("SHOT")
        r["_strats"] = ss
        for x in ss: st[x] += 1
    S["strategies"] = st.most_common(); S["shot_added_from_attachments"] = shot_auto
    S["no_strategy"] = sum(not r["_strats"] for r in C)
    S["gen_share"] = sum(r["tell"] == "GEN" for r in C) / len(C)
    per = {"jun_jul": lambda d: d < "2026-08", "aug_sep": lambda d: "2026-08" <= d < "2026-10", "oct_nov": lambda d: d >= "2026-10"}
    S["pole_by_period"] = {k: {"n": sum(f(r["date"]) for r in C),
                               "excess": sum(f(r["date"]) and r["tell"] in EX for r in C),
                               "bland": sum(f(r["date"]) and r["tell"] in BL for r in C),
                               "TYP": sum(f(r["date"]) and r["tell"] == "TYP" for r in C)} for k, f in per.items()}

    # monthly
    mon = collections.defaultdict(collections.Counter)
    for r in C: mon[r["date"][:7]][r["tell"]] += 1
    S["by_month"] = {m: sum(c.values()) for m, c in sorted(mon.items())}
    # recurrence within project
    th = collections.Counter((r["project"], r["tell"]) for r in C)
    S["threads"] = len(th); S["threads_recurring"] = sum(v >= 2 for v in th.values())
    S["corrections_in_recurring_threads"] = sum(v for v in th.values() if v >= 2)
    S["tells_in_3plus_projects"] = sorted([t for t, n in S["projects_per_tell"].items() if n >= 3])
    # law status at the time of each correction
    for r in C: r["_law"] = law_status(L, r)
    S["law_status"] = collections.Counter(r["_law"] for r in C)
    S["law_status_by_tell"] = {t: dict(collections.Counter(r["_law"] for r in C if r["tell"] == t)) for t in tells}
    S["law_status_by_month"] = {m: dict(collections.Counter(r["_law"] for r in C if r["date"][:7] == m)) for m in S["by_month"]}
    # before/after: per (project, tell) thread with a covering law, rate per active review day
    days = collections.defaultdict(set)
    for r in C: days[r["project"]].add(r["date"])
    bef_n = aft_n = bef_d = aft_d = 0; rows = []
    for (p, t), n in sorted(th.items()):
        ls = [l for l in L if l["scope"] in (p, "*") and t in l["tells"].split(";")]
        if not ls: continue
        ld = min(l["date"] for l in ls)
        rs = [r for r in C if r["project"] == p and r["tell"] == t]
        b = sum(r["date"] <= ld for r in rs); af = sum(r["date"] > ld for r in rs)  # same-day = the correction that made the law
        db = len([d for d in days[p] if d <= ld]); da = len([d for d in days[p] if d > ld])
        rows.append(dict(project=p, tell=t, law_date=ld, before=b, after=af, days_before=db, days_after=da))
        bef_n += b; aft_n += af; bef_d += db; aft_d += da
    S["law_threads"] = rows
    S["law_threads_n"] = len(rows)
    S["law_threads_after_zero"] = sum(x["after"] == 0 for x in rows)
    S["law_threads_after_zero_with_exposure"] = sum(x["after"] == 0 and x["days_after"] > 0 for x in rows)
    S["law_threads_with_exposure"] = sum(x["days_after"] > 0 for x in rows)
    S["rate_before"] = bef_n / bef_d if bef_d else None
    S["rate_after"] = aft_n / aft_d if aft_d else None
    S["before_after_totals"] = dict(before=bef_n, after=aft_n, days_before=bef_d, days_after=aft_d)
    S["owner_requested"] = sum(r["owner_requested"] == "1" for r in C)
    S["reversals"] = len(V); S["reversals_by_tell"] = collections.Counter(v["tell"] for v in V)
    S["quotes"] = sum(bool(r["quote"]) for r in C)
    # consistency check: seeded sample, independent keyword coder vs hand code
    rnd = random.Random(20261102); samp = rnd.sample(C, 60)
    kc = [kw_code(r["pattern"]) for r in samp]; hc = [r["tell"] for r in samp]
    po, k = kappa(hc, kc); S["consistency"] = dict(n=60, seed=20261102, agreement=round(po, 3), kappa=round(k, 3))
    with open(os.path.join(DATA, "consistency_check.csv"), "w", newline="") as f:
        w = csv.writer(f); w.writerow(["id", "pattern", "hand_tell", "keyword_tell", "agree"])
        for r, x in zip(samp, kc): w.writerow([r["id"], r["pattern"], r["tell"], x, int(x == r["tell"])])
    json.dump(S, open(os.path.join(DATA, "stats.json"), "w"), indent=1, default=list)
    # tables.md
    out = ["# Tables (generated by analyse.py)\n", "## Corrections by tell\n", "| Tell | Name | n | Projects | Law in force | Law elsewhere only | No law yet |", "|---|---|---|---|---|---|---|"]
    for t, n in S["by_tell"]:
        ls = S["law_status_by_tell"][t]
        out.append(f"| {t} | {tells[t]} | {n} | {S['projects_per_tell'][t]} | {ls.get('in_force',0)} | {ls.get('elsewhere',0)} | {ls.get('none',0)} |")
    out += ["\n## Corrections by project\n", "| Project | n |", "|---|---|"] + [f"| {p} | {n} |" for p, n in S["by_project"]]
    out += ["\n## Law threads (project x tell with a written rule)\n", "| Project | Tell | First rule | Before (incl. same day) | After | Review days before | after |", "|---|---|---|---|---|---|---|"]
    out += [f"| {x['project']} | {x['tell']} | {x['law_date']} | {x['before']} | {x['after']} | {x['days_before']} | {x['days_after']} |" for x in rows]
    out += ["\n## Strategies\n", "| Strategy | n |", "|---|---|"] + [f"| {k} | {v} |" for k, v in S["strategies"]]
    open(os.path.join(DATA, "tables.md"), "w").write("\n".join(out) + "\n")
    print(json.dumps({k: S[k] for k in ["n_corrections", "n_projects", "n_messages_coded", "pole", "law_status", "threads", "threads_recurring", "law_threads_n", "law_threads_after_zero", "law_threads_with_exposure", "law_threads_after_zero_with_exposure", "rate_before", "rate_after", "before_after_totals", "consistency", "strategies", "gen_share", "by_source", "memory_only", "multi_source", "tells_in_3plus_projects", "by_month", "law_status_by_month"]}, default=list, indent=0))
    if a.figdir:
        figures(C, S, tells, EX, BL, a.figdir)

def figures(C, S, tells, EX, BL, figdir):
    import matplotlib; matplotlib.use("svg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 11, "svg.fonttype": "none",
                         "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": "#555"})
    INK, MUTE = "#1d1d1f", "#6b6b73"
    PC = {"excess": "#c2410c", "bland": "#1f5f8b", "neither": "#8a8a93"}
    os.makedirs(figdir, exist_ok=True)
    def pole(t): return "excess" if t in EX else "bland" if t in BL else "neither"
    short = {"COL": "Default colour", "GLO": "Glow, glass, blur", "TYP": "Generic or robotic type",
             "ICO": "Generic icons, emoji", "IDN": "Generic marks", "DOT": "Status dots and pills",
             "SRF": "Surface defaults", "TPL": "Template composition", "CLU": "Clutter and copy tells",
             "MOT": "Generic motion", "DED": "Dead, flat, dull", "NAV": "Duplicate entry points",
             "CMP": "Stock components", "GIM": "Gimmicks", "GEN": "Unspecified 'AI feel'"}
    # Fig 1: tells
    items = sorted(S["by_tell"], key=lambda x: x[1])
    fig, ax = plt.subplots(figsize=(7.2, 5.4))
    ys = range(len(items))
    ax.barh(list(ys), [n for _, n in items], color=[PC[pole(t)] for t, _ in items], height=0.66)
    for y, (t, n) in zip(ys, items):
        ax.text(n + 0.4, y, f"{n}  ({S['projects_per_tell'][t]} projects)", va="center", fontsize=10.5, color=INK)
    ax.set_yticks(list(ys)); ax.set_yticklabels([short[t] for t, _ in items], fontsize=11)
    ax.set_xlim(0, max(n for _, n in items) * 1.45); ax.set_xlabel("Corrections", color=MUTE)
    from matplotlib.patches import Patch
    ax.legend(handles=[Patch(color=PC["excess"], label="Too much (added by default)"), Patch(color=PC["bland"], label="Too safe (generic by default)"), Patch(color=PC["neither"], label="Neither pole")], loc="lower right", frameon=False, fontsize=10)
    ax.set_title(f"{S['n_corrections']} design corrections by kind of AI-default tell", loc="left", fontsize=12.5, color=INK)
    fig.tight_layout(); fig.savefig(os.path.join(figdir, "fig1-tells.svg"), facecolor="white"); plt.close(fig)
    # Fig 2: law status by month
    months = list(S["by_month"]); lab = {"2026-06": "Jun", "2026-07": "Jul", "2026-08": "Aug", "2026-09": "Sep", "2026-10": "Oct", "2026-11": "1 Nov"}
    cats = [("none", "No written rule for it anywhere yet", "#b9b9c0"), ("elsewhere", "Rule existed in another project only", "#d97706"), ("in_force", "Rule already in force in this project", "#9f1239")]
    fig, ax = plt.subplots(figsize=(7.2, 4.4)); bottom = [0] * len(months)
    for k, name, col in cats:
        vals = [S["law_status_by_month"][m].get(k, 0) for m in months]
        ax.bar([lab.get(m, m) for m in months], vals, bottom=bottom, color=col, label=name, width=0.62)
        for i, v in enumerate(vals):
            if v >= 6: ax.text(i, bottom[i] + v / 2, str(v), ha="center", va="center", color="white", fontsize=10.5)
        bottom = [b + v for b, v in zip(bottom, vals)]
    ax.set_ylabel("Corrections", color=MUTE); ax.legend(frameon=False, fontsize=10, loc="upper left")
    ls = S["law_status"]
    ax.set_title(f"When each correction was made, was there already a written rule? ({ls.get('in_force',0)+ls.get('elsewhere',0)} of {S['n_corrections']}: yes)", loc="left", fontsize=11.5, color=INK)
    fig.tight_layout(); fig.savefig(os.path.join(figdir, "fig2-rule-status.svg"), facecolor="white"); plt.close(fig)
    # Fig 3: project x tell matrix
    projs = [p for p, _ in S["by_project"]]; tl = [t for t, _ in S["by_tell"]]
    cnt = collections.Counter((r["project"], r["tell"]) for r in C)
    pname = {"bupples": "Bupples", "wayclub": "WayClub", "portfolio": "Portfolio", "planket": "Planket", "partner-portfolio": "Partner portfolio", "onplan": "OnPlan", "adelante": "Adelante", "tajweed": "Tajweed", "nutrigo": "NutriGo", "experts-marine": "Experts Marine", "himalaya": "Himalaya concept", "spellquest": "Spell Quest", "hackfam": "HACKFAM", "photoshoot": "Photoshoot", "careeros": "CareerOS"}
    fig, ax = plt.subplots(figsize=(7.4, 6.0))
    mx = max(cnt.values())
    for i, p in enumerate(projs):
        for j, t in enumerate(tl):
            v = cnt.get((p, t), 0)
            if v:
                ax.scatter(j, i, s=40 + 520 * v / mx, color=PC[pole(t)], alpha=0.85, linewidths=0)
                if v >= 3: ax.text(j, i, str(v), ha="center", va="center", color="white", fontsize=9)
    ax.set_xticks(range(len(tl))); ax.set_xticklabels([short[t] for t in tl], rotation=55, ha="right", fontsize=10)
    ax.set_yticks(range(len(projs))); ax.set_yticklabels([pname.get(p, p) for p in projs], fontsize=10.5)
    ax.invert_yaxis(); ax.set_xlim(-0.6, len(tl) - 0.4); ax.grid(color="#eee", linewidth=0.6); ax.set_axisbelow(True)
    ax.set_title("Where each tell came up (dot area = corrections)", loc="left", fontsize=12, color=INK)
    fig.tight_layout(); fig.savefig(os.path.join(figdir, "fig3-matrix.svg"), facecolor="white"); plt.close(fig)
    # Fig 4: strategies
    sn = {"BAN": "Name the banned patterns", "BEN": "Benchmark a named product", "SHOT": "Screenshot or device review", "PRO": "Rendered options or prototype", "REF": "Own reference as the grammar", "RES": "Research mandate", "RESET": "Full visual reset", "AUD": "Audit another agent's work"}
    it = sorted(S["strategies"], key=lambda x: x[1])
    fig, ax = plt.subplots(figsize=(7.0, 3.9))
    ax.barh([sn[k] for k, _ in it], [v for _, v in it], color="#3f3f46", height=0.62)
    for y, (k, v) in enumerate(it): ax.text(v + 0.3, y, str(v), va="center", fontsize=10.5)
    ax.set_xlabel("Corrections that used it", color=MUTE)
    ax.set_title(f"How he steered: strategies visible in the {S['n_corrections']} corrections", loc="left", fontsize=12, color=INK)
    fig.tight_layout(); fig.savefig(os.path.join(figdir, "fig4-strategies.svg"), facecolor="white"); plt.close(fig)

if __name__ == "__main__":
    main()