"""OCR-like corruption of normalised transliterations.

Substitutions follow confusions between graphically or phonetically close runes and their
transliterations (k/g, t/þ/d, b/p, i/e, u/o/y, R/r, m/n, a/æ/ą, s/z); deletions and
insertions are drawn uniformly. `confusions` can be replaced by a matrix estimated from a
recogniser's own errors (see estimate_confusions).
"""
import random
from collections import Counter, defaultdict

DEFAULT_CONFUSIONS = {
    "k": "g", "g": "k", "t": "þd", "þ": "td", "d": "tþ", "b": "p", "p": "b",
    "i": "e", "e": "i", "u": "oy", "o": "u", "y": "u", "R": "r", "r": "R",
    "m": "n", "n": "m", "a": "æą", "æ": "a", "ą": "a", "s": "z", "z": "s",
    "l": "i", "f": "u", "h": "n",
}
LETTERS = "abdefghiklmnorstuþR"


def corrupt(text, cer, rng, confusions=None):
    """Return text with approximately `cer` character edits per reference character."""
    if cer <= 0 or not text:
        return text
    conf = confusions or DEFAULT_CONFUSIONS
    out = []
    for ch in text:
        if ch == " ":
            # word boundaries are dropped or kept, never substituted
            if rng.random() < cer * 0.3:
                continue
            out.append(ch)
            continue
        r = rng.random()
        if r < cer:
            op = rng.random()
            if op < 0.6:  # substitution
                cands = conf.get(ch)
                if isinstance(cands, dict):
                    keys, w = zip(*cands.items())
                    out.append(rng.choices(keys, weights=w)[0])
                else:
                    out.append(rng.choice(cands) if cands else rng.choice(LETTERS))
            elif op < 0.85:  # deletion
                continue
            else:  # insertion
                out.append(ch)
                out.append(rng.choice(LETTERS))
        else:
            out.append(ch)
    return "".join(out)


def estimate_confusions(pairs):
    """pairs: iterable of (reference, hypothesis). Returns {ref_char: {hyp_char: count}} from
    Levenshtein alignments (substitutions only)."""
    from rapidfuzz.distance import Levenshtein

    subs = defaultdict(Counter)
    for ref, hyp in pairs:
        for op, i, j in Levenshtein.editops(ref, hyp):
            if op == "replace" and ref[i] != " " and hyp[j] != " ":
                subs[ref[i]][hyp[j]] += 1
    return {k: dict(v) for k, v in subs.items()}


if __name__ == "__main__":
    rng = random.Random(0)
    s = "haraltr kunukR baþ kaurua kubl þausi aft kurm faþur sin"
    for c in (0.1, 0.3, 0.6):
        print(c, corrupt(s, c, rng))