"""Phase 2 — search + country vote + template reasoning. No LLM, fully offline.

Pure logic given an index (N,512 fp32, L2-normed), meta rows, and centroids.
The same module is imported by app.py (the Space) and eval/eval.py (the PC), so
the served algorithm and the evaluated algorithm are byte-identical.
"""

from dataclasses import dataclass

import numpy as np


@dataclass
class Config:
    K: int = 15          # top matches per view
    P: float = 8.0       # similarity sharpening exponent
    ALPHA: float = 0.4   # large-guide bias correction (S_c /= max(count[c], MIN_COUNT)**ALPHA)
    MIN_COUNT: int = 20  # floor on the per-country image count used for ALPHA normalization,
                         # so tiny guides (5–9 imgs) don't get a huge score boost and beat
                         # large countries on weak matches (the "India → mid-Pacific" bug)
    PRIOR_BETA: float = 0.3  # strength of the game-location prior (S_c *= available[c]**BETA).
                             # 0 disables the prior (filter still applies). Higher = trust the
                             # "how many real locations exist in this country" signal more.
    TEXT_WEIGHT: float = 0.25  # blend weight for the zero-shot TEXT branch (geo only):
                              # final = (1-W)*image_retrieval + W*image↔country_text. 0 disables.
    P_TEXT: float = 4.0       # sharpening exponent for zero-shot text sims
    JITTER_DEG: float = 0.0
    MAX_CLUES: int = 3
    RUNNER_SENTENCE_RATIO: float = 0.5  # show runner-up note when its score >= 50% of winner


# --- reasoning templates, keyed by confidence band; chosen deterministically ---
ASSERTIVE = [
    "Both views match features documented for {country} — {clue1} and {clue2}.{runner} Placing my guess in central {country}.",
    "This is {country} for me: I'm seeing {clue1}, and {clue2} backs it up.{runner} Dropping my pin in central {country}.",
    "Clear {country} signals here — {clue1} together with {clue2}.{runner} I'll guess central {country}.",
    "Confident on {country}: {clue1} and {clue2} both line up.{runner} Going central {country}.",
]
HEDGED = [
    "This looks most like {country} to me — I matched {clue1}, though I also saw similarities to {runner_c}. Going with {country}.",
    "Leaning {country}: {clue1} points that way, but {runner_c} crossed my mind too. I'll commit to {country}.",
    "Probably {country} — {clue1} is the strongest hint, with {runner_c} as a maybe. Guessing {country}.",
    "My read is {country} based on {clue1}, even if {runner_c} isn't far off. Placing it in {country}.",
]
UNCERTAIN = [
    "Tough one. Weak matches all around, but {clue1} nudges me toward {country}, so that's my guess.",
    "Not much to go on here — {clue1} is the only real hint, pointing at {country}. Rolling with it.",
    "Low confidence on this. {clue1} loosely suggests {country}, so I'll take the shot.",
    "Hard to read. {clue1} is faint, but it leans {country} — guessing there.",
]


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


def _distinct_clues(matches, k):
    """Up to k distinct clue texts, ordered by sim desc."""
    out, seen = [], set()
    for m in sorted(matches, key=lambda x: -x["sim"]):
        c = m["clue"]
        if c not in seen:
            seen.add(c)
            out.append(m)
        if len(out) >= k:
            break
    return out


def build_reasoning(country_name, confidence, winner_matches, runner_name, runner_matches, cfg):
    clues = _distinct_clues(winner_matches, cfg.MAX_CLUES)
    clue1 = clues[0]["clue"] if clues else "a few subtle details"
    clue2 = clues[1]["clue"] if len(clues) > 1 else clue1
    seed = "|".join(c["clue"] for c in clues) + country_name

    if confidence >= 0.55:
        runner = ""
        if runner_name and runner_matches:
            runner = f" I also weighed {runner_matches[0]['country_name']}, but those matches were weaker."
        return _pick(ASSERTIVE, seed).format(
            country=country_name, clue1=clue1.lower(), clue2=clue2.lower(), runner=runner)
    if confidence >= 0.35:
        runner_c = runner_matches[0]["country_name"] if runner_matches else "a neighbour"
        return _pick(HEDGED, seed).format(country=country_name, clue1=clue1.lower(), runner_c=runner_c)
    return _pick(UNCERTAIN, seed).format(country=country_name, clue1=clue1.lower())


def guess(embs, index, rows, centroids, count, cfg=Config(), priors=None,
          text_vecs=None, text_countries=None):
    """embs: list of (D,) float32 L2-normed view embeddings (1+; e.g. 5 frames around
    the spot + a downward car view). Returns the response dict.

    priors: optional {slug: available_location_count}. When given, the bot can ONLY
    guess countries present here (those the game actually has locations for) and
    weights each by available_count**PRIOR_BETA.

    text_vecs/(text_countries): optional (C,D) L2-normed country text embeddings and
    their aligned slugs (StreetCLIP zero-shot). When given, the final score blends
    image retrieval with image↔country-text similarity (cfg.TEXT_WEIGHT).
    """
    if not isinstance(embs, (list, tuple)):
        embs = [embs]
    matches = {}  # index row -> {sim, ...meta}
    for e in embs:
        sims = index @ e
        top = np.argpartition(-sims, cfg.K)[: cfg.K] if len(sims) > cfg.K else np.arange(len(sims))
        for idx in top:
            s = float(sims[idx])
            # if a row appears in both views, keep the larger sim but it still counts once here;
            # union semantics with double-contribution handled by summing weights per view below
            prev = matches.get(idx)
            if prev is None:
                r = rows[idx]
                matches[idx] = {"sim": s, "country": r["country"],
                                "country_name": r["country_name"], "clue": r["clue"],
                                "page": r["page"], "_w": max(s, 0.0) ** cfg.P}
            else:
                # second view also matched this row: add its weight (intended double vote)
                prev["_w"] += max(s, 0.0) ** cfg.P
                prev["sim"] = max(prev["sim"], s)

    # country scores
    scores, by_country = {}, {}
    for m in matches.values():
        c = m["country"]
        if priors is not None and priors.get(c, 0) <= 0:
            continue  # country has no game locations → the bot must never guess it
        scores[c] = scores.get(c, 0.0) + m["_w"]
        by_country.setdefault(c, []).append(m)
    for c in scores:
        scores[c] /= (max(count.get(c, 1), cfg.MIN_COUNT) ** cfg.ALPHA)
        if priors is not None:
            scores[c] *= priors[c] ** cfg.PRIOR_BETA

    if not scores:
        # None of the matched countries are in the game's pool — fall back to the
        # most location-rich country so we still return a valid, in-pool guess.
        winner = max(priors, key=priors.get) if priors else rows[0]["country"]
        cen = centroids[winner]
        return {
            "lat": cen["lat"], "lon": cen["lon"],
            "country": winner, "country_name": cen["name"],
            "confidence": 0.0, "runner_up": None,
            "reasoning": "Couldn't match these views to anywhere I know — taking a blind guess.",
            "clues": [],
        }

    # --- zero-shot TEXT branch (geo): blend image↔country-text with retrieval ---
    if text_vecs is not None and text_countries is not None and cfg.TEXT_WEIGHT > 0:
        sims_t = np.max(np.stack([text_vecs @ e for e in embs]), axis=0)  # best view per country
        zs = {}
        for i, c in enumerate(text_countries):
            if priors is not None and priors.get(c, 0) <= 0:
                continue
            s = max(float(sims_t[i]), 0.0) ** cfg.P_TEXT
            if priors is not None:
                s *= priors[c] ** cfg.PRIOR_BETA
            zs[c] = s
        rsum = sum(scores.values()) or 1.0
        zsum = sum(zs.values()) or 1.0
        W = cfg.TEXT_WEIGHT
        final = {
            c: (1 - W) * (scores.get(c, 0.0) / rsum) + W * (zs.get(c, 0.0) / zsum)
            for c in set(scores) | set(zs)
        }
    else:
        final = scores

    ranked = sorted(final.items(), key=lambda kv: -kv[1])
    winner, win_score = ranked[0]
    total = sum(final.values()) or 1.0
    confidence = win_score / total
    runner = ranked[1][0] if len(ranked) > 1 else None
    runner_score = ranked[1][1] if len(ranked) > 1 else 0.0

    win_matches = sorted(by_country.get(winner, []), key=lambda x: -x["sim"])
    runner_matches = sorted(by_country.get(runner, []), key=lambda x: -x["sim"]) if runner else []

    cen = centroids[winner]
    # The ratio gate only governs the optional extra runner-up SENTENCE in the
    # assertive band; the hedged band always names the actual runner-up.
    show_runner = bool(runner) and runner_score >= cfg.RUNNER_SENTENCE_RATIO * win_score
    if confidence >= 0.55:
        reasoning = build_reasoning(cen["name"], confidence, win_matches,
                                    runner if show_runner else None,
                                    runner_matches if show_runner else [], cfg)
    else:
        reasoning = build_reasoning(cen["name"], confidence, win_matches,
                                    runner, runner_matches, cfg)

    clue_list = []
    for m in _distinct_clues(win_matches, cfg.MAX_CLUES):
        clue_list.append({"text": m["clue"], "country": m["country"],
                          "page": m["page"], "sim": round(m["sim"], 3)})
    if runner_matches:
        m = runner_matches[0]
        clue_list.append({"text": m["clue"], "country": m["country"],
                          "page": m["page"], "sim": round(m["sim"], 3)})

    return {
        "lat": cen["lat"], "lon": cen["lon"],
        "country": winner, "country_name": cen["name"],
        "confidence": round(confidence, 3),
        "runner_up": runner,
        "reasoning": reasoning,
        "clues": clue_list,
    }