"""A5. Currier A/B, subject sections, and Davis hands: who predicts vocabulary?
- keyness A vs B (log-likelihood G2)
- per-page B-score from diagnostic features, plotted in codex order
- cosine similarity between (section, language, hand) groups + clustering
Outputs: out/tables/keyness_AB.csv, out/tables/group_similarity.csv,
         out/figs/bscore_pages.pdf, out/figs/group_similarity.pdf
"""
import os, sys
_here = os.path.dirname(os.path.abspath(__file__))  # data/analysis, or a subfolder of it
sys.path.insert(0, _here if os.path.exists(os.path.join(_here, "paths.py")) else os.path.dirname(_here))
from paths import *  # noqa: repository locations
import collections, csv, math, os
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from corpus import Corpus, section_name
import ivtff

C = Corpus()
A = C.running_text(L={"A"}); B = C.running_text(L={"B"})
cA, cB = collections.Counter(A), collections.Counter(B); nA, nB = len(A), len(B)


def g2(a, b, na, nb):
    e1 = na * (a + b) / (na + nb); e2 = nb * (a + b) / (na + nb)
    g = 0
    if a: g += a * math.log(a / e1)
    if b: g += b * math.log(b / e2)
    return 2 * g * (1 if a / na > b / nb else -1)


key = []
for w in set(cA) | set(cB):
    if cA[w] + cB[w] < 20: continue
    key.append((w, cA[w], cB[w], cA[w] / nA * 1e4, cB[w] / nB * 1e4, g2(cA[w], cB[w], nA, nB)))
key.sort(key=lambda r: r[5])
with open("out/tables/keyness_AB.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["word", "count_A", "count_B", "per10k_A", "per10k_B", "G2_signed"])
    for r in key[:25] + key[-25:]: w.writerow(r)
print("B-key:", [(r[0], int(r[5])) for r in key[:15]])
print("A-key:", [(r[0], int(r[5])) for r in key[-15:]])

# --- per-page B-score: share of word tokens containing 'ed' or ending 'edy'/'eedy' (B features)
# minus share of 'chol/chor/cho*' type (A features); use simple, transparent features.
feat_B = lambda w: ("ed" in w) or w.endswith("eey") or w.startswith("qok")
feat_A = lambda w: w.startswith("cho") or w.startswith("sho") or w == "daiin" or w.endswith("ol") or w.endswith("or")
pages = []
for p in C.pages:
    ws = C.running_text(pages={p})
    if len(ws) < 30: continue
    b = sum(map(feat_B, ws)) / len(ws); a = sum(map(feat_A, ws)) / len(ws)
    v = C.pagevars[p]
    pages.append(dict(page=p, folio=ivtff.folio_number(p), section=section_name(v.get("I")), lang=v.get("L", "-"), hand=v.get("H", "-"), n=len(ws), fB=b, fA=a, score=b - a))
with open("out/tables/bscore_pages.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=pages[0].keys()); w.writeheader(); w.writerows(pages)
fig, ax = plt.subplots(figsize=(7.0, 3.0))
cols = {"H": "#4c9a4c", "B": "#2aa6a6", "P": "#d49a3a", "S": "#c85c5c", "T": "#999999", "C": "#8a74c9", "A": "#5d7fc9", "Z": "#3a4f8a"}
x = np.arange(len(pages))
for i, r in enumerate(pages):
    I = C.pagevars[r["page"]].get("I")
    ax.bar(i, r["score"], color=cols.get(I, "#999"), width=1.0, edgecolor="none")
    ax.plot(i, -0.62, marker="s", ms=2.2, color={"A": "black", "B": "#c0392b"}.get(r["lang"], "#bbbbbb"), ls="none")
ax.axhline(0, color="black", lw=0.5)
ax.set_xlim(-1, len(pages)); ax.set_ylim(-0.7, 0.65)
ax.set_ylabel("B-score"); ax.set_xlabel("pages in current binding order (pages with ≥30 running-text words)")
# folio tick labels every ~10
ticks = [i for i, r in enumerate(pages) if r["folio"] % 10 == 0 and r["page"].endswith("r")]
ax.set_xticks(ticks); ax.set_xticklabels([f"f{pages[i]['folio']}" for i in ticks], fontsize=7)
import matplotlib.patches as mpatches
ax.legend(handles=[mpatches.Patch(color=cols[k], label=section_name(k)) for k in ["H", "P", "B", "S", "C", "T"]] +
          [plt.Line2D([], [], marker="s", ls="none", color="black", label="Currier A"), plt.Line2D([], [], marker="s", ls="none", color="#c0392b", label="Currier B")],
          fontsize=6, frameon=False, ncol=4, loc="upper left")
ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); fig.savefig("out/figs/bscore_pages.pdf")

# --- group similarity
groups = collections.OrderedDict()
for p in C.pages:
    v = C.pagevars[p]; I = v.get("I"); L = v.get("L", "-"); H = v.get("H", "-")
    if I in ("Z",): continue
    k = f"{section_name(I)} / {L} / hand {H}"
    groups.setdefault(k, []).extend(C.running_text(pages={p}))
groups = {k: v for k, v in groups.items() if len(v) >= 500}
vocab = [w for w, c in collections.Counter(w for v in groups.values() for w in v).most_common(400)]
M = np.array([[collections.Counter(v)[w] / len(v) for w in vocab] for v in groups.values()])
M = M / np.linalg.norm(M, axis=1, keepdims=True)
S = M @ M.T
names = list(groups.keys())
with open("out/tables/group_similarity.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["group", "tokens"] + names)
    for i, n in enumerate(names): w.writerow([n, len(groups[n])] + [f"{s:.3f}" for s in S[i]])
print(open("out/tables/group_similarity.csv").read())
fig, ax = plt.subplots(figsize=(6.2, 5.2))
im = ax.imshow(S, cmap="Greys", vmin=0.3, vmax=1.0)
ax.set_xticks(range(len(names))); ax.set_yticks(range(len(names)))
ax.set_xticklabels([n + f"  (n={len(groups[n])})" for n in names], rotation=90, fontsize=7); ax.set_yticklabels(names, fontsize=7)
for i in range(len(names)):
    for j in range(len(names)):
        ax.text(j, i, f"{S[i, j]:.2f}", ha="center", va="center", fontsize=6, color="white" if S[i, j] > 0.7 else "black")
fig.tight_layout(); fig.savefig("out/figs/group_similarity.pdf")