"""A2. Core statistics of the running text vs. reference corpora of equal size.
Zipf, type/token & hapax, word length, character entropy (h0,h1,h2) on the
EVA alphabet and on a glyph-merged alphabet, index of coincidence.
Outputs: out/tables/basic_stats.csv, out/figs/zipf.pdf, wordlen.pdf, entropy.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 math, csv, collections, os
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from corpus import Corpus, tokenize_glyphs
import refcorpora

os.makedirs("out/tables", exist_ok=True); os.makedirs("out/figs", exist_ok=True)
C = Corpus()
V_all = C.running_text()
V_A = C.running_text(L={"A"})
V_B = C.running_text(L={"B"})
N = len(V_all)
refs = refcorpora.all_refs(N)
print("Voynich running text tokens:", N, "A:", len(V_A), "B:", len(V_B))


def entropies(words, unit=lambda w: list(w)):
    """h0 = log2(alphabet), h1 = unigram entropy, h2 = conditional entropy H(X_n|X_{n-1}),
    computed on the character stream with word boundaries as a symbol."""
    seq = []
    for w in words:
        seq += unit(w) + [" "]
    uni = collections.Counter(seq); n = len(seq)
    h1 = -sum(c / n * math.log2(c / n) for c in uni.values())
    bi = collections.Counter(zip(seq, seq[1:])); nb = len(seq) - 1
    h2 = 0.0
    for (a, b), c in bi.items():
        p_ab = c / nb; p_b_given_a = c / uni[a]
        h2 -= p_ab * math.log2(p_b_given_a)
    ic = sum(c * (c - 1) for s, c in uni.items() if s != " ") / (sum(c for s, c in uni.items() if s != " ") * (sum(c for s, c in uni.items() if s != " ") - 1))
    return dict(alphabet=len(uni) - 1, h0=math.log2(len(uni) - 1), h1=h1, h2=h2, ic=ic)


def basic(words):
    n = len(words); types = collections.Counter(words)
    hapax = sum(1 for c in types.values() if c == 1)
    lens = [len(w) for w in words]
    mean = sum(lens) / n; var = sum((l - mean) ** 2 for l in lens) / n
    # Zipf slope: least squares on log rank vs log freq for ranks 1..1000
    freqs = sorted(types.values(), reverse=True)[:1000]
    xs = [math.log(i + 1) for i in range(len(freqs))]; ys = [math.log(f) for f in freqs]
    mx, my = sum(xs) / len(xs), sum(ys) / len(ys)
    slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sum((x - mx) ** 2 for x in xs)
    top = types.most_common(1)[0]
    return dict(tokens=n, types=len(types), ttr=len(types) / n, hapax=hapax, hapax_frac=hapax / len(types),
                mean_len=mean, sd_len=math.sqrt(var), zipf_slope=slope, top_word=top[0], top_share=top[1] / n)


rows = []
corpora = {"Voynich (all running text)": V_all, "Voynich A": V_A, "Voynich B": V_B}
corpora.update(refs)
for name, ws in corpora.items():
    b = basic(ws); e = entropies(ws)
    r = dict(corpus=name, **b, **{k: v for k, v in e.items()})
    if name.startswith("Voynich"):
        eg = entropies(ws, tokenize_glyphs)
        r.update(glyph_alphabet=eg["alphabet"], glyph_h1=eg["h1"], glyph_h2=eg["h2"], glyph_ic=eg["ic"])
    rows.append(r)
keys = sorted({k for r in rows for k in r}, key=lambda k: (k != "corpus", k))
with open("out/tables/basic_stats.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()})

# --- Zipf figure
fig, ax = plt.subplots(figsize=(5.2, 3.6))
styles = {"Voynich (all running text)": dict(color="black", lw=2.0)}
for name, ws in corpora.items():
    if name in ("Voynich A", "Voynich B"): continue
    freqs = sorted(collections.Counter(ws).values(), reverse=True)
    ax.loglog(range(1, len(freqs) + 1), freqs, label=name, **styles.get(name, dict(lw=1.0, alpha=0.8)))
ax.set_xlabel("rank"); ax.set_ylabel("frequency"); ax.legend(fontsize=7, frameon=False)
ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); fig.savefig("out/figs/zipf.pdf")

# --- word length figure
fig, ax = plt.subplots(figsize=(5.2, 3.4))
for name, ws in corpora.items():
    if name in ("Voynich A", "Voynich B"): continue
    c = collections.Counter(len(w) for w in ws); n = len(ws)
    xs = list(range(1, 16)); ax.plot(xs, [c[x] / n for x in xs], marker="o", ms=3, label=name, **styles.get(name, dict(lw=1.0, alpha=0.8)))
ax.set_xlabel("word length (characters; EVA for Voynich)"); ax.set_ylabel("share of tokens")
ax.legend(fontsize=7, frameon=False); ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); fig.savefig("out/figs/wordlen.pdf")

# --- entropy figure (h1 vs h2)
fig, ax = plt.subplots(figsize=(5.2, 3.6))
for r in rows:
    if r["corpus"] in ("Voynich A", "Voynich B"): continue
    ax.scatter(r["h1"], r["h2"], color="black" if r["corpus"].startswith("Voynich") else "#777777", s=28)
    ax.annotate(r["corpus"].replace(" (all running text)", " (EVA)"), (r["h1"], r["h2"]), fontsize=7, xytext=(4, 3), textcoords="offset points")
    if "glyph_h1" in r:
        ax.scatter(r["glyph_h1"], r["glyph_h2"], color="black", s=28, marker="s")
        ax.annotate("Voynich (glyph units)", (r["glyph_h1"], r["glyph_h2"]), fontsize=7, xytext=(4, -9), textcoords="offset points")
ax.set_xlabel("h1: single-character entropy (bits)"); ax.set_ylabel("h2: conditional entropy (bits)")
ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout(); fig.savefig("out/figs/entropy.pdf")

# --- top 30 words table
top = collections.Counter(V_all).most_common(30)
cA, cB = collections.Counter(V_A), collections.Counter(V_B)
with open("out/tables/top_words.csv", "w", newline="") as f:
    w = csv.writer(f); w.writerow(["word", "count", "share", "per_10k_A", "per_10k_B"])
    for word, c in top:
        w.writerow([word, c, f"{c / N:.4f}", f"{cA[word] / len(V_A) * 1e4:.1f}", f"{cB[word] / len(V_B) * 1e4:.1f}"])
print(open("out/tables/top_words.csv").read())