"""Normalisation of model answers into one comparable record per question.

Pure Python (math only): no torch, no model code, so it can be tested anywhere.
Moved verbatim from app/models.py in 1.2.0.
"""
from __future__ import annotations

import math


def option_keys(q: dict) -> list[str]:
    qtype = q.get("type", "choice")
    crit = q.get("criteria", q.get("options"))
    if qtype == "noul":
        return ["true", "false"]
    if qtype == "score":
        return [str(i) for i in range(len(crit or []))]
    if isinstance(crit, dict):
        return [str(k) for k in crit]
    return [str(c) for c in (crit or [])]


def _lookup(raw: dict, key: str):
    """Tolerant key lookup: exact, str(), lower-case, and bool spellings."""
    if not isinstance(raw, dict):
        return None
    candidates = [key, key.lower(), key.capitalize()]
    if key == "true":
        candidates += [True, "True", "yes", "Yes", 1, "1"]
    if key == "false":
        candidates += [False, "False", "no", "No", 0, "0"]
    if key.isdigit():
        candidates.append(int(key))
    for c in candidates:
        if c in raw:
            return raw[c]
    return None


def normalise(q: dict, probs_raw: dict | None, choice=None, p_true=None, level=None) -> dict:
    """One comparable record per question, whatever the model returned."""
    qtype = q.get("type", "choice")
    keys = option_keys(q)
    probs = {}
    if qtype == "noul" and p_true is not None:
        p = float(p_true)
        probs = {"true": p, "false": 1.0 - p}
    elif probs_raw:
        for k in keys:
            v = _lookup(probs_raw, k)
            probs[k] = float(v) if v is not None else 0.0
    if not probs or sum(probs.values()) <= 0:
        # The model gave only its answer: represent it as a point mass so the UI still works.
        probs = {k: 0.0 for k in keys}
        if qtype == "score" and level is not None:
            probs[str(int(round(float(level))))] = 1.0
        elif choice is not None and str(choice).lower() in {k.lower() for k in keys}:
            probs[next(k for k in keys if k.lower() == str(choice).lower())] = 1.0
    s = sum(probs.values()) or 1.0
    probs = {k: v / s for k, v in probs.items()}
    top = max(probs, key=probs.get)
    n = len(probs)
    ent = -sum(p * math.log(p) for p in probs.values() if p > 0)
    ent_conf = 1.0 - ent / math.log(n) if n > 1 else 1.0
    ent_conf = max(0.0, min(1.0, ent_conf))
    # Laya reports 1 - normalised entropy for choice/score, and max(p, 1 - p) for yes/no questions.
    laya_conf = probs[top] if qtype == "noul" else ent_conf
    rec = dict(type=qtype, choice=top, probs=probs, top_prob=probs[top], entropy_conf=ent_conf, laya_conf=laya_conf)
    if qtype == "score":
        rec["expected_level"] = sum(int(k) * p for k, p in probs.items())
        rec["levels"] = len(keys)
    if qtype == "noul":
        rec["p_true"] = probs["true"]
    return rec