"""Corpusul contemporan, în două registre care se citesc separat.

  scris   — Wikipedia în română, textul enciclopedic de azi
  vorbit  — replici de film subtitrate în română (OpenSubtitles 2018)

Nu le adunăm într-o singură cifră. Un cuvânt turcesc și un cuvânt francez au
sorți diferite după registru, iar diferența e una dintre concluziile paginii.

Ieșire: raw/now_counts.json
"""
import bz2, json, os, re, sys
from collections import Counter

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from lemma import Lemmatizer, normalize, tokenize

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "..", "raw")

# Cât text de Wikipedia citim. Enciclopedia are peste 3 GB decomprimați; un
# eșantion de articole luate în ordinea dumpului e destul pentru frecvențe.
WIKI_CHAR_TARGET = int(os.environ.get("WIKI_CHARS", 120_000_000))

TEMPLATE_RE = re.compile(r"\{\{[^{}]*\}\}")


def strip_wikitext(s):
    s = s.replace("&lt;", "<").replace("&gt;", ">").replace("&amp;", "&")
    s = re.sub(r"<ref[^>]*/>", " ", s)
    s = re.sub(r"<ref[^>]*>.*?</ref>", " ", s, flags=re.S)
    s = re.sub(r"<!--.*?-->", " ", s, flags=re.S)
    s = re.sub(r"<[^>]{1,80}>", " ", s)
    for _ in range(6):
        new = TEMPLATE_RE.sub(" ", s)
        if new == s:
            break
        s = new
    s = re.sub(r"\{\|.*?\|\}", " ", s, flags=re.S)
    s = re.sub(r"\[\[(?:Fișier|File|Imagine|Categorie|Category):[^\]]*\]\]", " ", s)
    s = re.sub(r"\[\[[^\]|]*\|", " ", s)
    s = re.sub(r"\[\[|\]\]", " ", s)
    s = re.sub(r"https?://\S+", " ", s)
    s = re.sub(r"^[*#:;=|!].*$", " ", s, flags=re.M)
    s = re.sub(r"'{2,}", "", s)
    return s


def wikipedia_counts():
    src = os.path.join(RAW, "rowiki.xml.bz2")
    counts = Counter()
    chars = arts = 0
    buf, in_text, in_page, ns0 = [], False, False, False
    with bz2.open(src, "rt", encoding="utf8", errors="replace") as fh:
        for line in fh:
            if "<page>" in line:
                in_page, ns0, buf = True, False, []
                continue
            if not in_page:
                continue
            if "<ns>0</ns>" in line:
                ns0 = True
            if "<redirect" in line:
                in_page = False
                continue
            if "<text" in line:
                in_text = True
                line = line[line.index("<text"):]
                line = line[line.index(">") + 1:] if ">" in line else ""
            if in_text:
                if "</text>" in line:
                    buf.append(line[:line.index("</text>")])
                    in_text = False
                    if ns0:
                        body = strip_wikitext("".join(buf))
                        counts.update(tokenize(body))
                        chars += len(body)
                        arts += 1
                        if arts % 40000 == 0:
                            print(f"    {arts:,} articole, {chars:,} caractere", flush=True)
                    in_page = False
                    buf = []
                else:
                    buf.append(line)
            if chars >= WIKI_CHAR_TARGET:
                break
    print(f"  Wikipedia: {arts:,} articole, {chars:,} caractere, "
          f"{sum(counts.values()):,} cuvinte")
    return counts, {"articles": arts, "chars": chars}


def subtitle_counts():
    counts = Counter()
    path = os.path.join(RAW, "ro_full.txt")
    lines = 0
    with open(path, encoding="utf8") as f:
        for line in f:
            p = line.split()
            if len(p) != 2:
                continue
            w, c = p[0], int(p[1])
            if not re.fullmatch(r"[a-zA-ZăâîșțĂÂÎȘȚşţŞŢ'-]+", w):
                continue
            counts[normalize(w)] += c
            lines += 1
    print(f"  Subtitrări: {lines:,} tipuri de formă, {sum(counts.values()):,} cuvinte")
    return counts, {"form_types": lines}


def main():
    L = Lemmatizer()
    out = {}
    print("subtitrări...")
    sc, smeta = subtitle_counts()
    slem, sunk = L.resolve(sc)
    out["vorbit"] = dict(
        smeta, tokens=sum(sc.values()), unmatched_tokens=sunk,
        lemma_tokens=round(sum(slem.values())),
        lemmas={k: round(v, 2) for k, v in slem.items() if v >= 2},
    )
    print("Wikipedia...")
    wc, wmeta = wikipedia_counts()
    wlem, wunk = L.resolve(wc)
    out["scris"] = dict(
        wmeta, tokens=sum(wc.values()), unmatched_tokens=wunk,
        lemma_tokens=round(sum(wlem.values())),
        lemmas={k: round(v, 2) for k, v in wlem.items() if v >= 2},
    )
    for k, v in out.items():
        print(f"  {k}: {v['lemma_tokens']:,} cuvinte lematizate, "
              f"nepotrivite {v['unmatched_tokens']/max(v['tokens'],1)*100:.1f}%")
    with open(os.path.join(RAW, "now_counts.json"), "w", encoding="utf8") as f:
        json.dump(out, f, ensure_ascii=False)
    print("scris raw/now_counts.json")


if __name__ == "__main__":
    main()
