"""A4. The line and the paragraph as units: line-initial / line-final glyph
distributions vs. line-internal; paragraph-first-line gallows enrichment;
word length by position in line.
Outputs: out/tables/line_effects.csv, out/figs/line_effects.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, re, os
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from corpus import Corpus
import ivtff

C = Corpus()
P = [l for l in C.loci_where(generic="P")]

first_init, last_init, mid_init = collections.Counter(), collections.Counter(), collections.Counter()
first_fin, last_fin, mid_fin = collections.Counter(), collections.Counter(), collections.Counter()
lens_by_pos = collections.defaultdict(list)
par_first_gallows, par_other_gallows = [0, 0], [0, 0]
nlines = 0
for l in P:
    ws = [w for w in ivtff.words(l.text) if "?" not in w]
    if len(ws) < 3: continue
    nlines += 1
    first_init[ws[0][0]] += 1; last_init[ws[-1][0]] += 1
    first_fin[ws[0][-1]] += 1; last_fin[ws[-1][-1]] += 1
    for w in ws[1:-1]:
        mid_init[w[0]] += 1; mid_fin[w[-1]] += 1
    for i, w in enumerate(ws):
        key = "first" if i == 0 else ("last" if i == len(ws) - 1 else "middle")
        lens_by_pos[key].append(len(w))
    # paragraph-initial lines: locator '@' or '*' with type P0/P1 (new paragraph)
    is_par_start = l.par_start
    g = sum(bool(re.match(r"^[ktpf]|^c[ktpf]h", w)) for w in ws) / len(ws)
    # first word of line starting with gallows
    g1 = bool(re.match(r"^[ktpf]|^c[ktpf]h", ws[0]))
    tgt = par_first_gallows if is_par_start else par_other_gallows
    tgt[0] += g1; tgt[1] += 1


def share(c, keys):
    n = sum(c.values()); return {k: c[k] / n for k in keys}


rows = []
for name, cnt in [("line-first word, initial glyph", first_init), ("line-middle word, initial glyph", mid_init), ("line-last word, initial glyph", last_init)]:
    rows.append(dict(position=name, **share(cnt, list("ocqsdyapktfl"))))
for name, cnt in [("line-first word, final glyph", first_fin), ("line-middle word, final glyph", mid_fin), ("line-last word, final glyph", last_fin)]:
    rows.append(dict(position=name, **share(cnt, list("ynlrsmodg"))))
keys = sorted({k for r in rows for k in r}, key=lambda k: (k != "position", k))
with open("out/tables/line_effects.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows)
for r in rows: print({k: (round(v, 3) if isinstance(v, float) else v) for k, v in r.items()})
print("lines used:", nlines)
print("mean word length first/middle/last:", {k: round(sum(v) / len(v), 2) for k, v in lens_by_pos.items()})
print("paragraph-first lines beginning with a gallows: %.3f (n=%d); other lines: %.3f (n=%d)" % (par_first_gallows[0] / par_first_gallows[1], par_first_gallows[1], par_other_gallows[0] / par_other_gallows[1], par_other_gallows[1]))
with open("out/tables/line_effects.csv", "a", newline="") as f:
    w = csv.writer(f)
    w.writerow([]); w.writerow(["lines_used", nlines])
    for k, v in lens_by_pos.items(): w.writerow([f"mean_len_{k}", sum(v) / len(v)])
    w.writerow(["par_first_line_gallows_initial", par_first_gallows[0] / par_first_gallows[1], par_first_gallows[1]])
    w.writerow(["other_line_gallows_initial", par_other_gallows[0] / par_other_gallows[1], par_other_gallows[1]])

# figure
fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.8))
for ax, (keys, trip, title) in zip(axes, [(list("ocqsdypkt"), (first_init, mid_init, last_init), "initial glyph of word"),
                                          (list("ynlrsmod"), (first_fin, mid_fin, last_fin), "final glyph of word")]):
    import numpy as np
    x = np.arange(len(keys)); wdt = 0.27
    for i, (cnt, lab, col) in enumerate(zip(trip, ["first in line", "middle", "last in line"], ["#222222", "#888888", "#cccccc"])):
        s = share(cnt, keys); ax.bar(x + (i - 1) * wdt, [s[k] for k in keys], wdt, color=col, edgecolor="#222222", linewidth=0.3, label=lab)
    ax.set_xticks(x); ax.set_xticklabels(keys); ax.set_title(title, fontsize=9); ax.spines[["top", "right"]].set_visible(False)
axes[0].legend(fontsize=6.5, frameon=False); axes[0].set_ylabel("share of words")
fig.tight_layout(); fig.savefig("out/figs/line_effects.pdf")