"""Atlas — the LEARNED model's inference. A linear classifier head on top of the
frozen StreetCLIP embedding: probs = softmax(emb @ W + b) over playable countries,
averaged across the round's frames, then weighted by the game-location prior.

Unlike guess.py (nearest-neighbour retrieval), this is a trained model: it learned
what distinguishes each country from all PlonkIt + geohints images.
"""

import numpy as np

ASSERTIVE = [
    "I've learned to read scenes like this as {country} — the overall mix of road, signage and surroundings fits.",
    "This reads clearly as {country} to me, weighing everything in view.",
    "Confident on {country}: the combination of cues I trained on lines up.",
    "My read is {country} — the whole scene matches what I know for it.",
]
HEDGED = [
    "Looks most like {country}, though {runner} crossed my mind. Going with {country}.",
    "Leaning {country} here; {runner} was the next best. Committing to {country}.",
    "Probably {country} — {runner} was close, but I'll take {country}.",
    "My best read is {country}, with {runner} as a maybe.",
]
UNCERTAIN = [
    "Tough scene — nothing jumps out, but it leans {country}, so that's my guess.",
    "Low confidence, but the overall feel points to {country}.",
    "Hard to read; I'll take {country} as the most likely.",
    "Not sure, but {country} fits best, so guessing there.",
]


def load_head(path):
    z = np.load(path, allow_pickle=True)
    return (z["W"].astype(np.float32), z["b"].astype(np.float32),
            [str(c) for c in z["classes"]])


def _pick(variants, seed_text):
    return variants[abs(hash(seed_text)) % len(variants)]


SCRIPT_PENALTY = 0.15   # multiply prob of countries that don't use the detected (non-Latin) script
# Only these distinctive scripts are reliable zero-shot; sinhala/lao/bengali/etc. are
# noisy attractors that misfire on text-less scenes, so they're excluded from detection.
RELIABLE_SCRIPTS = {"thai", "cyrillic", "greek", "cjk", "devanagari", "arabic", "hebrew"}
SCRIPT_MARGIN = 0.015   # the detected script must beat Latin by at least this


def _detect_script(embs, script_vecs, script_names):
    """Return the detected writing system, or None (Latin / text-less / uncertain).
    Fires only when a RELIABLE non-Latin script clearly beats Latin."""
    if "latin" not in script_names:
        return None
    sims = np.zeros(len(script_names))
    for e in embs:
        sims += script_vecs @ e
    sims /= len(embs)
    latin = sims[script_names.index("latin")]
    best, best_s = None, -1e9
    for i, name in enumerate(script_names):
        if name in RELIABLE_SCRIPTS and sims[i] > best_s:
            best, best_s = name, sims[i]
    return best if best is not None and (best_s - latin) >= SCRIPT_MARGIN else None


def predict(embs, W, b, classes, centroids, cfg, priors=None,
            script_vecs=None, script_names=None, country_scripts=None):
    """embs: list of (768,) StreetCLIP embeddings. Returns the response dict."""
    if not isinstance(embs, (list, tuple)):
        embs = [embs]
    probs = np.zeros(len(classes), dtype=np.float64)
    for e in embs:
        logits = e @ W + b
        logits = logits - logits.max()
        ex = np.exp(logits)
        probs += ex / ex.sum()
    probs /= len(embs)

    # weight by the game-location prior (and hard-filter 0-location countries)
    if priors is not None:
        for i, c in enumerate(classes):
            av = priors.get(c, 0)
            probs[i] = 0.0 if av <= 0 else probs[i] * (av ** cfg.PRIOR_BETA)
        s = probs.sum()
        if s > 0:
            probs /= s

    # script branch: if a distinctive (non-Latin) writing system is detected,
    # down-weight countries that don't use it (Thai → only Thailand, etc.)
    detected = None
    if script_vecs is not None and country_scripts is not None:
        detected = _detect_script(embs, script_vecs, script_names)
        if detected:
            for i, c in enumerate(classes):
                if detected not in country_scripts.get(c, ["latin"]):
                    probs[i] *= SCRIPT_PENALTY
            s = probs.sum()
            if s > 0:
                probs /= s

    order = np.argsort(-probs)
    winner = classes[order[0]]
    confidence = float(probs[order[0]])
    runner = classes[order[1]] if len(order) > 1 else None
    cen = centroids[winner]
    runner_name = centroids.get(runner, {}).get("name", runner) if runner else None

    seed = winner + f"{confidence:.2f}"
    if confidence >= 0.55:
        reasoning = _pick(ASSERTIVE, seed).format(country=cen["name"])
    elif confidence >= 0.30:
        reasoning = _pick(HEDGED, seed).format(country=cen["name"], runner=runner_name or "a neighbour")
    else:
        reasoning = _pick(UNCERTAIN, seed).format(country=cen["name"])
    if detected:
        reasoning += f" (I can see {detected} script, which points here.)"

    return {
        "lat": cen["lat"], "lon": cen["lon"],
        "country": winner, "country_name": cen["name"],
        "confidence": round(confidence, 3),
        "runner_up": runner,
        "runner_up_name": runner_name,
        "runner_up_conf": round(float(probs[order[1]]), 3) if len(order) > 1 else None,
        "reasoning": reasoning,
        "clues": [],
        "script": detected,
    }