from __future__ import annotations

"""Modern-language word popularity (Zipf + within-language rank).

Primary: Robyn Speer's ``wordfreq`` (Wikipedia, subtitles, news, books, web, social).
Fallback: HermitDave FrequencyWords (OpenSubtitles) for languages wordfreq lacks —
practical stand-in for popular usage when FastText crawl ``.bin`` models (multi-GB
each) cannot ship in a Space image. Optional ``FASTTEXT_FREQ_DIR/{lang}.txt`` ranked
vocab files are also consulted when present (one word per line, most frequent first).
"""

import math
import os
import re
from functools import lru_cache
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
FREQ_DIR = ROOT / "data" / "freq"
FASTTEXT_DIR = Path(os.environ.get("FASTTEXT_FREQ_DIR", str(FREQ_DIR / "fasttext")))

# Atlas lang key / ISO 639-3 → wordfreq / FrequencyWords codes.
LANG_TO_CODE: dict[str, str] = {
    "english": "en",
    "spanish": "es",
    "french": "fr",
    "german": "de",
    "italian": "it",
    "portuguese": "pt",
    "dutch": "nl",
    "russian": "ru",
    "polish": "pl",
    "swedish": "sv",
    "norwegian": "nb",
    "danish": "da",
    "finnish": "fi",
    "hungarian": "hu",
    "czech": "cs",
    "slovak": "sk",
    "romanian": "ro",
    "bulgarian": "bg",
    "greek": "el",
    "turkish": "tr",
    "arabic": "ar",
    "hebrew": "he",
    "hindi": "hi",
    "bengali": "bn",
    "indonesian": "id",
    "malay": "ms",
    "vietnamese": "vi",
    "thai": "th",
    "chinese": "zh",
    "japanese": "ja",
    "korean": "ko",
    "ukrainian": "uk",
    "catalan": "ca",
    "croatian": "hr",
    "serbian": "sr",
    "slovenian": "sl",
    "lithuanian": "lt",
    "latvian": "lv",
    "estonian": "et",
    "persian": "fa",
    "urdu": "ur",
    "tamil": "ta",
    "tagalog": "tl",
    "filipino": "tl",
    "icelandic": "is",
    "basque": "eu",
    "galician": "gl",
    "eng": "en",
    "spa": "es",
    "fra": "fr",
    "fre": "fr",
    "deu": "de",
    "ger": "de",
    "ita": "it",
    "por": "pt",
    "nld": "nl",
    "dut": "nl",
    "rus": "ru",
    "pol": "pl",
    "swe": "sv",
    "nor": "nb",
    "nob": "nb",
    "dan": "da",
    "fin": "fi",
    "hun": "hu",
    "ces": "cs",
    "cze": "cs",
    "slk": "sk",
    "ron": "ro",
    "rum": "ro",
    "bul": "bg",
    "ell": "el",
    "gre": "el",
    "tur": "tr",
    "arb": "ar",
    "heb": "he",
    "hin": "hi",
    "ben": "bn",
    "ind": "id",
    "msa": "ms",
    "vie": "vi",
    "tha": "th",
    "cmn": "zh",
    "zho": "zh",
    "jpn": "ja",
    "kor": "ko",
    "ukr": "uk",
    "cat": "ca",
    "hrv": "hr",
    "srp": "sr",
    "slv": "sl",
    "lit": "lt",
    "lav": "lv",
    "est": "et",
    "fas": "fa",
    "pes": "fa",
    "urd": "ur",
    "tam": "ta",
    "tgl": "tl",
    "fil": "tl",
    "isl": "is",
    "eus": "eu",
    "glg": "gl",
}

# FrequencyWords filename stem on HermitDave/FrequencyWords content/2018/
FW_FILES: dict[str, str] = {
    "en": "en_50k.txt",
    "es": "es_50k.txt",
    "fr": "fr_50k.txt",
    "de": "de_50k.txt",
    "it": "it_50k.txt",
    "pt": "pt_50k.txt",
    "nl": "nl_50k.txt",
    "ru": "ru_50k.txt",
    "pl": "pl_50k.txt",
    "sv": "sv_50k.txt",
    "cs": "cs_50k.txt",
    "ro": "ro_50k.txt",
    "hu": "hu_50k.txt",
    "tr": "tr_50k.txt",
    "uk": "uk_50k.txt",
    "fi": "fi_50k.txt",
    "da": "da_50k.txt",
    "el": "el_50k.txt",
    "bg": "bg_50k.txt",
    "hr": "hr_50k.txt",
    "sk": "sk_50k.txt",
    "nb": "no_50k.txt",
    "ca": "ca_50k.txt",
    "id": "id_50k.txt",
    "vi": "vi_50k.txt",
    "ar": "ar_50k.txt",
    "he": "he_50k.txt",
    "hi": "hi_50k.txt",
    "ko": "ko_50k.txt",
    "ja": "ja_50k.txt",
    "zh": "zh_50k.txt",
    "fa": "fa_50k.txt",
    "tl": "tl_50k.txt",
}

FW_BASE = "https://raw.githubusercontent.com/hermitdave/FrequencyWords/master/content/2018/"

_WORD_RE = re.compile(r"[^\W\d_]+", re.UNICODE)


def freq_code_for(lang: str, iso_639_3: str | None = None) -> str | None:
    """Map atlas language → frequency code. Only known modern codes (no bare 2-letter guess)."""
    key = (lang or "").strip().casefold()
    if key in LANG_TO_CODE:
        return LANG_TO_CODE[key]
    iso = (iso_639_3 or "").strip().casefold()
    if iso in LANG_TO_CODE:
        return LANG_TO_CODE[iso]
    return None


def _norm_term(term: str) -> str:
    t = (term or "").strip().casefold()
    if t.startswith("*"):
        t = t[1:]
    return t


@lru_cache(maxsize=64)
def _wordfreq_ranks(code: str) -> dict[str, tuple[float, int]] | None:
    try:
        from wordfreq import get_frequency_dict, available_languages
    except ImportError:
        return None
    langs = available_languages(wordlist="best")
    if code not in langs:
        langs = available_languages(wordlist="large")
        wordlist = "large" if code in langs else None
        if wordlist is None:
            return None
    else:
        wordlist = "best"
    freq_dict = get_frequency_dict(code, wordlist=wordlist)
    # Zipf = log10(occurrences per billion words) == log10(proportion) + 9
    ranked = sorted(freq_dict.items(), key=lambda kv: -kv[1])
    out: dict[str, tuple[float, int]] = {}
    for i, (w, f) in enumerate(ranked, start=1):
        if f <= 0:
            continue
        z = math.log10(f) + 9.0
        out[w] = (z, i)
        if i >= 1_000_000:
            break
    return out


@lru_cache(maxsize=64)
def _ranked_file_map(path: str) -> dict[str, tuple[float, int]]:
    """Ranked word list → (pseudo-zipf, rank). Zipf ≈ 7.5 - log10(rank)."""
    out: dict[str, tuple[float, int]] = {}
    p = Path(path)
    if not p.exists():
        return out
    with p.open(encoding="utf-8", errors="ignore") as fh:
        rank = 0
        for line in fh:
            line = line.strip()
            if not line:
                continue
            # FrequencyWords: "word count" ; fasttext vocab dump: "word" or "word freq"
            parts = line.split()
            if not parts:
                continue
            w = parts[0].casefold()
            if not w or w.startswith("#"):
                continue
            rank += 1
            zipf = max(0.0, 7.5 - math.log10(rank))
            out.setdefault(w, (zipf, rank))
            if rank >= 1_000_000:
                break
    return out


def ensure_frequencywords(code: str) -> Path | None:
    fname = FW_FILES.get(code)
    if not fname:
        return None
    FREQ_DIR.mkdir(parents=True, exist_ok=True)
    dest = FREQ_DIR / fname
    if dest.exists() and dest.stat().st_size > 0:
        return dest
    url = FW_BASE + fname
    try:
        import httpx

        with httpx.Client(timeout=60.0, follow_redirects=True) as client:
            r = client.get(url)
            r.raise_for_status()
            dest.write_bytes(r.content)
        return dest
    except Exception:
        return None


def _fasttext_path(code: str) -> Path | None:
    for name in (f"{code}.txt", f"{code}_freq.txt", f"cc.{code}.txt"):
        p = FASTTEXT_DIR / name
        if p.exists():
            return p
    return None


def _lookup_in_map(mp: dict[str, tuple[float, int]], norm: str) -> tuple[float, int] | None:
    """Exact form first; single-token fallback only (never first word of a phrase)."""
    hit = mp.get(norm)
    if hit is not None:
        return hit
    # Avoid "Cape Verde" → "cape" false popularity.
    if any(ch.isspace() for ch in norm) or "-" in norm:
        return None
    m = _WORD_RE.fullmatch(norm) or _WORD_RE.search(norm)
    if m and m.group(0) != norm:
        return mp.get(m.group(0))
    return None


def lookup_frequency(term: str, lang: str, iso_639_3: str | None = None) -> dict:
    """Return zipf, rank (1=most frequent), and source. Unknown → zipf 0, rank 0."""
    code = freq_code_for(lang, iso_639_3)
    norm = _norm_term(term)
    empty = {"zipf": 0.0, "rank": 0, "source": None, "code": code}
    if not code or not norm:
        return empty

    # 1) wordfreq
    wf = _wordfreq_ranks(code)
    if wf is not None:
        hit = _lookup_in_map(wf, norm)
        if hit is not None:
            return {"zipf": hit[0], "rank": hit[1], "source": "wordfreq", "code": code}

    # 2) optional FastText ranked vocab dir
    ft = _fasttext_path(code)
    if ft is not None:
        mp = _ranked_file_map(str(ft))
        hit = _lookup_in_map(mp, norm)
        if hit is not None:
            return {"zipf": hit[0], "rank": hit[1], "source": "fasttext", "code": code}

    # 3) FrequencyWords (OpenSubtitles) popular-usage fallback
    fw = ensure_frequencywords(code)
    if fw is not None:
        mp = _ranked_file_map(str(fw))
        hit = _lookup_in_map(mp, norm)
        if hit is not None:
            return {"zipf": hit[0], "rank": hit[1], "source": "frequencywords", "code": code}

    return empty


def _wordfreq_hits_for(code: str, needed: set[str]) -> dict[str, tuple[float, int]] | None:
    """Rank + Zipf for ``needed`` forms only (avoids multi-100k dicts in Docker builds)."""
    if not needed:
        return {}
    try:
        from wordfreq import get_frequency_dict, available_languages
    except ImportError:
        return None
    langs = available_languages(wordlist="best")
    if code not in langs:
        langs = available_languages(wordlist="large")
        wordlist = "large" if code in langs else None
        if wordlist is None:
            return None
    else:
        wordlist = "best"
    freq_dict = get_frequency_dict(code, wordlist=wordlist)
    remaining = set(needed)
    out: dict[str, tuple[float, int]] = {}
    # Iterate in descending frequency so enumerate position is the rank.
    for i, (w, f) in enumerate(sorted(freq_dict.items(), key=lambda kv: -kv[1]), start=1):
        if w in remaining and f > 0:
            out[w] = (math.log10(f) + 9.0, i)
            remaining.discard(w)
            if not remaining:
                break
        if i >= 1_000_000:
            break
    return out


def _file_hits_for(path: Path, needed: set[str]) -> dict[str, tuple[float, int]]:
    if not needed or not path.exists():
        return {}
    out: dict[str, tuple[float, int]] = {}
    remaining = set(needed)
    with path.open(encoding="utf-8", errors="ignore") as fh:
        rank = 0
        for line in fh:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if not parts:
                continue
            w = parts[0].casefold()
            if not w or w.startswith("#"):
                continue
            rank += 1
            if w in remaining:
                out[w] = (max(0.0, 7.5 - math.log10(rank)), rank)
                remaining.discard(w)
                if not remaining:
                    break
            if rank >= 1_000_000:
                break
    return out


def _map_for_needed(code: str, needed: set[str]) -> tuple[dict[str, tuple[float, int]], str | None]:
    hits = _wordfreq_hits_for(code, needed)
    if hits is not None:
        return hits, "wordfreq"
    ft = _fasttext_path(code)
    if ft is not None:
        return _file_hits_for(ft, needed), "fasttext"
    fw = ensure_frequencywords(code)
    if fw is not None:
        return _file_hits_for(fw, needed), "frequencywords"
    return {}, None


def _map_for_code(code: str) -> tuple[dict[str, tuple[float, int]], str | None]:
    """Load the best available rank map for a language code once (runtime lookups)."""
    wf = _wordfreq_ranks(code)
    if wf is not None:
        return wf, "wordfreq"
    ft = _fasttext_path(code)
    if ft is not None:
        return _ranked_file_map(str(ft)), "fasttext"
    fw = ensure_frequencywords(code)
    if fw is not None:
        return _ranked_file_map(str(fw)), "frequencywords"
    return {}, None


def _needed_forms(terms: list[str], idxs: list[int]) -> set[str]:
    needed: set[str] = set()
    for i in idxs:
        norm = _norm_term(terms[i])
        if not norm:
            continue
        needed.add(norm)
        if any(ch.isspace() for ch in norm) or "-" in norm:
            continue
        m = _WORD_RE.search(norm)
        if m and m.group(0) != norm:
            needed.add(m.group(0))
    return needed


def annotate_nodes(
    terms: list[str],
    langs: list[str],
    iso_by_lang: dict[str, str | None],
) -> tuple["object", "object"]:
    """Build per-node zipf (float32) and rank (uint32) arrays."""
    import numpy as np

    n = len(terms)
    zipf = np.zeros(n, dtype=np.float32)
    rank = np.zeros(n, dtype=np.uint32)
    by_code: dict[str, list[int]] = {}
    for i, lang in enumerate(langs):
        code = freq_code_for(lang, iso_by_lang.get(lang))
        if not code:
            continue
        by_code.setdefault(code, []).append(i)

    for code, idxs in sorted(by_code.items(), key=lambda kv: -len(kv[1])):
        needed = _needed_forms(terms, idxs)
        mp, source = _map_for_needed(code, needed)
        print(f"  {code}: {len(idxs):,} nodes via {source or 'none'} ({len(mp):,} hits / {len(needed):,} forms)")
        if mp:
            for i in idxs:
                norm = _norm_term(terms[i])
                if not norm:
                    continue
                hit = _lookup_in_map(mp, norm)
                if hit is not None:
                    zipf[i] = hit[0]
                    rank[i] = hit[1]
        _wordfreq_ranks.cache_clear()
        _ranked_file_map.cache_clear()
        del mp, needed
    return zipf, rank