#!/usr/bin/env python3
"""
Alinea el texto de la repo (lo auditado) contra el OCR del escaneo (la verdad)
y saca candidatos a correccion, clasificados por el consenso de los DOS tesseract.

Para cada libro:
  - repo  = concatenacion normalizada del markdown del libro
  - scan6 = concatenacion normalizada del OCR psm6
  - scan4 = concatenacion normalizada del OCR psm4
  difflib(repo, scan6) y difflib(repo, scan4) -> ops 'replace'.
  Un candidato:
    CONFIRMADO  si scan6 y scan4 leen lo MISMO (ambos tesseract coinciden y
                difieren de la repo) -> el escaneo gana, la repo se corrige.
    DUDOSO      si scan6 y scan4 discrepan entre si en ese punto -> escalar a
                Surya + Llama.
  Se ignora ruido de paginacion: inserciones solo-escaneo, y replaces donde
  ''.join(repo)==''.join(scan) (guiones/espaciado).

Uso: python3 diff_align.py [--book BOOKID] > out/<book>.candidates.jsonl
     python3 diff_align.py --all     (todos los libros con OCR disponible)
Salida JSONL: un candidato por linea.
"""
import json, re, sys, unicodedata, argparse, difflib
from pathlib import Path

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "books.json").read_text())
REPO_ROOT = Path(CFG["_repo_root"])
OCR_DIR = HERE / "ocr"
OUT = HERE / "out"; OUT.mkdir(exist_ok=True)

# --- normalizacion compartida ---
def strip_diacritics(s):
    return "".join(c for c in unicodedata.normalize("NFKD", s)
                   if not unicodedata.combining(c))

WORD_RE = re.compile(r"[a-z0-9]+")
def tokenize(text):
    """Devuelve lista de (token, raw_offset) normalizada para alinear."""
    text = strip_diacritics(text).lower()
    text = text.replace("’", "'").replace("‘", "'")
    toks = []
    for m in WORD_RE.finditer(text):
        toks.append((m.group(0), m.start()))
    return toks

# --- texto de la repo (solo ingles: traduccion + significado + glosas) ---
def repo_text(book):
    bdir = REPO_ROOT / book
    parts = []
    files = sorted(bdir.rglob("*.md"))
    for f in files:
        raw = f.read_text(errors="ignore")
        # quitar bloques de verso devanagari/transliteracion (blockquotes con '>')
        lines = []
        for ln in raw.splitlines():
            s = ln.strip()
            if s.startswith(">"):          # sloka / transliteracion: fuera
                continue
            if s.startswith("#"):          # encabezados de verso: fuera del flujo
                continue
            lines.append(ln)
        body = "\n".join(lines)
        # quitar markup
        body = re.sub(r"\*\*(.+?)\*\*", r"\1", body)
        body = re.sub(r"\*(.+?)\*", r"\1", body)
        body = re.sub(r"`([^`]*)`", r"\1", body)
        body = re.sub(r"<[^>]+>", " ", body)
        parts.append((str(f.relative_to(REPO_ROOT)), body))
    return parts

def repo_tokens(book):
    toks = []           # (token, fileref, offset)
    for fileref, body in repo_text(book):
        for tok, off in tokenize(body):
            toks.append((tok, fileref, off))
    return toks

# --- texto del escaneo OCR (concatena paginas en orden de PDF y pagina) ---
def scan_tokens(book, psm):
    bdir = OCR_DIR / book
    if not bdir.exists():
        return None
    toks = []           # (token, pageref)
    for stemdir in sorted(bdir.iterdir()):
        if not stemdir.is_dir():
            continue
        for pf in sorted(stemdir.glob(f"p*.psm{psm}.txt")):
            pageref = f"{stemdir.name}#{pf.name.split('.')[0]}"
            txt = pf.read_text(errors="ignore")
            for tok, off in tokenize(txt):
                toks.append((tok, pageref))
    return toks

def words(toks):
    return [t[0] for t in toks]

# --- diccionario ingles para distinguir error real de ruido OCR ---
def load_dict():
    d = set()
    p = Path("/usr/share/dict/words")
    if p.exists():
        for w in p.read_text(errors="ignore").splitlines():
            d.add(w.strip().lower())
    return d
ENGLISH = load_dict()

def known_set(toks):
    return set(t[0] for t in toks)

def is_known(tok, vocab):
    return len(tok) <= 2 or tok in ENGLISH or tok in vocab

# pares de caracteres que tesseract confunde tipicamente
CONFUSE = {frozenset(p) for p in [
    "ec","tl","oa","nu","il","ha","mn","sf","ft","ie","uo","co","vy",
    "rn","cl","li","tf","ea","ou","au","hb","ji","yg","gq","08","l1","i1"]}
def _lev_ops(a, b):
    sm = difflib.SequenceMatcher(None, a, b, autojunk=False)
    return sm.get_opcodes()
def looks_like_ocr(repo_str, scan_str):
    """True si es un solo token cuya diferencia son confusiones OCR clasicas."""
    rt, st = repo_str.split(), scan_str.split()
    if len(rt) != 1 or len(st) != 1:
        return False
    a, b = rt[0], st[0]
    if a == b:
        return False
    bad = 0
    for tag, i1, i2, j1, j2 in _lev_ops(a, b):
        if tag == "equal":
            continue
        if tag == "replace" and (i2 - i1) == 1 and (j2 - j1) == 1:
            if frozenset(a[i1] + b[j1]) in CONFUSE:
                continue
            bad += 1
        elif tag in ("insert", "delete") and (max(i2 - i1, j2 - j1) == 1):
            bad += 1   # un caracter de mas/menos: tolerable solo si poco
        else:
            bad += 2
    return bad == 0

def classify(repo_str, scan_str, vocab, consensus):
    rt = repo_str.split()
    st = [t for t in scan_str.split() if not t.isdigit()]
    repo_known = all(is_known(t, vocab) for t in rt) if rt else False
    scan_known = all(is_known(t, vocab) for t in st) if st else False
    if repo_known and not scan_known:
        return "OCR_NOISE"          # repo correcto, escaneo basura OCR
    if (not repo_known) and scan_known:
        return "REPO_TYPO"          # repo erroneo, escaneo es palabra real
    if repo_known and scan_known:
        if looks_like_ocr(repo_str, scan_str):
            return "OCR_NOISE"      # palabra real pero diferencia = confusion OCR
        return "REVISAR"            # ambas palabras reales, difieren de verdad
    return "AMBIGUO"

def join(seq):
    return "".join(seq)

CONTEXT = 6
def candidates_for(book):
    rtoks = repo_tokens(book)
    s6 = scan_tokens(book, 6)
    s4 = scan_tokens(book, 4)
    if not rtoks or not s6 or not s4:
        return []
    rw, w6, w4 = words(rtoks), words(s6), words(s4)

    def replaces(rwords, swords):
        """dict: (i1,i2) span en repo -> (repo_str, scan_str, j1,j2)
        Solo ediciones LOCALES y cortas. Para libros grandes divide en ventanas
        de CHUNK_R palabras para evitar O(n^2) de difflib en secuencias largas."""
        ANCHOR = 3
        CHUNK_R = 10000  # palabras repo por ventana difflib
        S_MARGIN = 2000  # palabras extra en la ventana de escaneo a cada lado
        EDGE = ANCHOR + 2  # suprimir candidatos cerca del borde de la ventana

        def _window(rchunk, schunk, r_off, s_off):
            sm = difflib.SequenceMatcher(None, rchunk, schunk, autojunk=False)
            ops = sm.get_opcodes()
            res = {}
            for k, (tag, i1, i2, j1, j2) in enumerate(ops):
                if tag != "replace":
                    continue
                if (i2 - i1) > 4 or (j2 - j1) > 6:
                    continue
                prev_ok = k > 0 and ops[k-1][0] == "equal" and (ops[k-1][2]-ops[k-1][1]) >= ANCHOR
                next_ok = k < len(ops)-1 and ops[k+1][0] == "equal" and (ops[k+1][2]-ops[k+1][1]) >= ANCHOR
                if not (prev_ok and next_ok):
                    continue
                rp = " ".join(rchunk[i1:i2])
                sp = " ".join(schunk[j1:j2])
                if join(rchunk[i1:i2]) == join(schunk[j1:j2]):
                    continue
                res[(r_off + i1, r_off + i2)] = (rp, sp, s_off + j1, s_off + j2)
            return res

        out = {}
        r_len = len(rwords)
        s_len = len(swords)
        ratio = s_len / r_len if r_len else 1.0
        step = max(CHUNK_R - 2 * EDGE, 1)
        r_start = 0
        while r_start < r_len:
            r_end = min(r_start + CHUNK_R, r_len)
            s_start = max(0, int(r_start * ratio) - S_MARGIN)
            s_end = min(s_len, int(r_end * ratio) + S_MARGIN)
            chunk_res = _window(rwords[r_start:r_end], swords[s_start:s_end], r_start, s_start)
            for (i1, i2), v in chunk_res.items():
                loc_i1 = i1 - r_start
                loc_i2 = i2 - r_start
                # suppress near left boundary (except first chunk): prev window owns these
                if r_start > 0 and loc_i2 < EDGE:
                    continue
                # suppress near right boundary (except last chunk): next window owns these
                if r_end < r_len and loc_i1 >= (r_end - r_start) - EDGE:
                    continue
                if (i1, i2) not in out:
                    out[(i1, i2)] = v
            r_start += step
        return out

    r6 = replaces(rw, w6)
    r4 = replaces(rw, w4)
    vocab = known_set(rtoks)

    cands = []
    spans = set(r6) | set(r4)
    for span in sorted(spans):
        i1, i2 = span
        rp6 = r6.get(span); rp4 = r4.get(span)
        repo_str = (rp6 or rp4)[0]
        scan6 = rp6[1] if rp6 else "(=repo)"
        scan4 = rp4[1] if rp4 else "(=repo)"
        # consenso entre los dos tesseract
        consensus = bool(rp6 and rp4 and rp6[1] == rp4[1])
        scan_str = rp6[1] if rp6 else rp4[1]
        # clasificacion lexica (diccionario + vocab repo)
        verdict = classify(repo_str, scan_str, vocab, consensus)
        # contexto repo
        lo, hi = max(0, i1 - CONTEXT), min(len(rw), i2 + CONTEXT)
        ctx = " ".join(rw[lo:i1]) + " [[" + repo_str + "]] " + " ".join(rw[i2:hi])
        fileref = rtoks[i1][1] if i1 < len(rtoks) else ""
        # pagina del escaneo (de scan6 si existe)
        pageref = ""
        if rp6:
            j = rp6[2]
            if j < len(s6): pageref = s6[j][1]
        elif rp4:
            j = rp4[2]
            if j < len(s4): pageref = s4[j][1]
        cands.append({
            "book": book, "file": fileref, "page": pageref,
            "verdict": verdict, "consensus": consensus,
            "repo": repo_str, "scan6": scan6, "scan4": scan4,
            "scan": scan_str, "context": ctx,
            "span": [i1, i2],
        })
    return cands

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--book")
    ap.add_argument("--all", action="store_true")
    a = ap.parse_args()
    books = [a.book] if a.book else list(CFG["books"].keys())
    grand = 0
    for book in books:
        if not (OCR_DIR / book).exists():
            if a.book: print(f"sin OCR: {book}", file=sys.stderr)
            continue
        cands = candidates_for(book)
        from collections import Counter
        cc = Counter(c["verdict"] for c in cands)
        outf = OUT / f"{book}.candidates.jsonl"
        with outf.open("w") as fh:
            for c in cands:
                fh.write(json.dumps(c, ensure_ascii=False) + "\n")
        print(f"{book}: {len(cands)} cand  "
              f"REPO_TYPO={cc['REPO_TYPO']} REVISAR={cc['REVISAR']} "
              f"AMBIGUO={cc['AMBIGUO']} OCR_NOISE={cc['OCR_NOISE']} -> {outf.name}",
              file=sys.stderr)
        grand += len(cands)
    print(f"TOTAL candidatos: {grand}", file=sys.stderr)

if __name__ == "__main__":
    main()