"""T7. Unsupervised HMM over the glyph stream (Baum–Welch, 2 and 3 states): do the
glyphs partition into hidden classes the way consonants and vowels do (Knight & Yamada
1999; Goldsmith & Xanthos 2009)? Run on Voynichese and on Latin, Italian, Hebrew (abjad)
and Malagasy; report the class each symbol is assigned to and the class-transition matrix.
Outputs: out/hmm_classes.csv
"""
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 csv, collections, sys
import numpy as np
from harness import *


def baum_welch(seq, K, S, iters=40, seed=0):
    rng = np.random.RandomState(seed); T = len(seq)
    A = rng.dirichlet(np.ones(S), S); B = rng.dirichlet(np.ones(K), S); pi = np.ones(S) / S
    for it in range(iters):
        # forward-backward in log-free scaled form
        alpha = np.zeros((T, S)); c = np.zeros(T)
        alpha[0] = pi * B[:, seq[0]]; c[0] = alpha[0].sum(); alpha[0] /= c[0]
        for t in range(1, T):
            alpha[t] = (alpha[t - 1] @ A) * B[:, seq[t]]; c[t] = alpha[t].sum(); alpha[t] /= c[t]
        beta = np.zeros((T, S)); beta[-1] = 1
        for t in range(T - 2, -1, -1):
            beta[t] = (A @ (B[:, seq[t + 1]] * beta[t + 1])) / c[t + 1]
        gamma = alpha * beta; gamma /= gamma.sum(1, keepdims=True)
        xi = np.zeros((S, S))
        for t in range(T - 1):
            x = alpha[t][:, None] * A * (B[:, seq[t + 1]] * beta[t + 1])[None, :]; xi += x / x.sum()
        A = xi / xi.sum(1, keepdims=True)
        Bn = np.zeros((S, K))
        for k in range(K): Bn[:, k] = gamma[seq == k].sum(0)
        B = (Bn + 1e-3) / (Bn + 1e-3).sum(1, keepdims=True); pi = gamma[0]
        ll = np.log(c).sum()
    return A, B, ll


def classes(words, units, label, S=2):
    seq_sym = []
    for w in words: seq_sym += (tokenize_glyphs(w) if units else list(w))
    cnt = collections.Counter(seq_sym); alpha = [a for a, c in cnt.most_common() if c >= 100]
    idx = {a: i for i, a in enumerate(alpha)}
    seq = np.array([idx[a] for a in seq_sym if a in idx])[:120000]
    A, B, ll = baum_welch(seq, len(alpha), S)
    # assign each symbol to argmax state of P(state|symbol) ∝ B[s,k]*pi_s(stationary)
    w, v = np.linalg.eig(A.T); stat = np.real(v[:, np.argmax(np.real(w))]); stat /= stat.sum()
    post = B * stat[:, None]; post /= post.sum(0, keepdims=True)
    rows = []
    for k, a in enumerate(alpha):
        rows.append(dict(corpus=label, symbol=a, count=cnt[a], state=int(post[:, k].argmax()), p=round(float(post[:, k].max()), 2)))
    return rows, A, ll


out = []
V = flatten(structured())
for label, words, units in [("Voynich", V, True)] + [(L, languages()[L][:len(V)], False) for L in ["Latin", "Italian", "Hebrew", "Malagasy", "Arabic"]]:
    for S in (2, 3):
        rows, A, ll = classes(words, units, label, S)
        for r in rows: r["states"] = S; out.append(r)
        by = collections.defaultdict(list)
        for r in rows: by[r["state"]].append(r["symbol"])
        print(f"{label} S={S}: " + " | ".join(f"{s}: {' '.join(v)}" for s, v in sorted(by.items())))
        print("   A=", np.round(A, 2).tolist())
with open("out/hmm_classes.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=out[0].keys()); w.writeheader(); w.writerows(out)