"""Which words a user's prompts keep dwelling on, for choosing where a camera move aims.

Each candidate target already gets a score from its own prompt (an owned part beats a bare
object, objects of a verb beat doers). This adds a second: how many DIFFERENT prompts the word
turned up in. "It keeps coming back, so it matters" is fused into the first score as a bonus.
Silent until MIN_PROMPTS prompts have been seen: one prompt's words are not a habit. A prompt
seen again counts nothing, so regenerating does not inflate a word.

Self-contained: no dependency on the host pack. One store per user, at
`<ComfyUI user dir>/shot_camera/focus_memory.json` (else `~/.shot_camera/focus_memory.json`);
the SHOT_CAMERA_MEMORY environment variable names another file. Delete the file to forget.
"""

import json
import math
import os

MIN_PROMPTS = 3
FREQ_W = 0.6            # max bonus: reorders near-equals, never outranks an owned part
MAX_HASHES = 300
PICK_W = 0.9            # bonus a word the user has chosen grows toward (tanh of picks)
REJECT_W = 0.6          # penalty a word the rewriter proposed and the user replaced grows toward
MIN_SHOTS_FOR_RATE = 10 # shots reviewed before "how often the user wants a move" is trusted
WORD_RATE_W = 0.5       # max bonus/penalty a word earns from the ratings of runs that aimed at it


def _path():
    override = os.environ.get("SHOT_CAMERA_MEMORY")
    if override:
        return override
    try:
        import folder_paths
        base = folder_paths.get_user_directory()
    except Exception:  # noqa: BLE001 — not inside ComfyUI, or an older one
        base = os.path.join(os.path.expanduser("~"), ".shot_camera")
        return os.path.join(base, "focus_memory.json")
    return os.path.join(base, "shot_camera", "focus_memory.json")


def _read():
    try:
        with open(_path(), "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except (OSError, ValueError):
        return {}


def prior():
    """{lemma: bonus}: words that keep coming back across prompts (once enough prompts were
    seen), plus what the user's own choices said: up for words they picked over the rewriter's
    proposal, down for the words they replaced."""
    data = _read()
    seen = data.get("seen") or {}
    out = {}
    if int(data.get("prompts", 0)) >= MIN_PROMPTS and seen:
        top = math.log1p(max(seen.values()))
        out = {lemma: FREQ_W * math.log1p(n) / top for lemma, n in seen.items()}
    for lemma, n in (data.get("picks") or {}).items():
        out[lemma] = out.get(lemma, 0.0) + PICK_W * math.tanh(n / 2.0)
    for lemma, n in (data.get("rejects") or {}).items():
        out[lemma] = out.get(lemma, 0.0) - REJECT_W * math.tanh(n / 2.0)
    for key, (g, b) in arm_stats().items():       # what ratings taught about aiming at a word
        if key.startswith("word:"):
            lemma = key[5:]
            out[lemma] = out.get(lemma, 0.0) + WORD_RATE_W * 2.0 * (arm_rate(key) - 0.5)
    return out


def learn(decisions):
    """Remember what a person chose. `decisions`: [{"auto": lemma|None, "picked": lemma|None,
    "mode": "auto"|"move"|"hold"|"none"}], one per shot shown. A picked word is reinforced; the
    word the rewriter proposed instead is marked down when they differ; "none" counts against
    moves in general (see move_rate). -> shots remembered."""
    data = _read()
    picks, rejects = data.get("picks") or {}, data.get("rejects") or {}
    total, kept = int(data.get("shots", 0)), int(data.get("kept", 0))
    n = 0
    for d in decisions or []:
        mode = d.get("mode") or "auto"
        n += 1
        total += 1
        kept += mode != "none"
        if mode == "none":
            continue
        picked, auto = d.get("picked"), d.get("auto")
        picked = [picked] if isinstance(picked, str) else list(picked or [])
        for word in picked:
            picks[word] = picks.get(word, 0) + 1
        if picked and auto and auto not in picked:
            rejects[auto] = rejects.get(auto, 0) + 1
    data.update(picks=picks, rejects=rejects, shots=total, kept=kept)
    _save(data)
    return n


def view_stats():
    """{key: (good, bad)} for each view and each "view@trait" the ratings and picks taught."""
    return {k: (float(v[0]), float(v[1])) for k, v in (_read().get("views") or {}).items()
            if isinstance(v, (list, tuple)) and len(v) == 2}


def _bump_views(data, used, good, bad):
    views = data.setdefault("views", {})
    for u in used:
        for key in [u["view"]] + [f"{u['view']}@{t}" for t in (u.get("traits") or ["none"])]:
            g, b = views.get(key, [0.0, 0.0])
            views[key] = [round(g + good, 4), round(b + bad, 4)]


def rate_views(used, sign):
    """A rated generation: its views share the credit (sign > 0) or the blame (sign < 0), so a
    run with four views does not punish each as hard as a run with one. `used`: [{"view",
    "traits"}]. -> views rated."""
    used = [u for u in used or [] if isinstance(u, dict) and u.get("view")]
    if not used or not sign:
        return 0
    data = _read()
    share = 1.0 / len(used)
    _bump_views(data, used, share if sign > 0 else 0.0, share if sign < 0 else 0.0)
    _save(data)
    return len(used)


def learn_views(decisions):
    """What a person chose: [{"auto": view|None, "picked": view|None, "traits": [...]}]. A pick
    counts as a good sign for that view; the view the rewriter proposed instead, as half a bad."""
    data = _read()
    n = 0
    for d in decisions or []:
        tr = d.get("traits") or []
        picked = [d["picked"]] if isinstance(d.get("picked"), str) else list(d.get("picked") or [])
        if picked:
            for v in picked:
                _bump_views(data, [{"view": v, "traits": tr}], 1.0 / len(picked), 0.0)
            n += 1
            if d.get("auto") and d["auto"] not in picked:
                _bump_views(data, [{"view": d["auto"], "traits": tr}], 0.0, 0.5)
    _save(data)
    return n


def arm_stats():
    """{arm: (good, bad)} for what ratings taught about the camera's own choices: "move:yes" /
    "move:no" (a shot got a move / was left alone), "style:hold|travel|k1|k2|k3", "word:<lemma>",
    "split:yes|no" (a shot whose point changes was cut in two / kept whole)."""
    return {k: (float(v[0]), float(v[1])) for k, v in (_read().get("arms") or {}).items()
            if isinstance(v, (list, tuple)) and len(v) == 2}


def arm_rate(key, stats=None):
    """Smoothed good-rate of one arm: 0.5 until rated."""
    g, b = (stats if stats is not None else arm_stats()).get(key, (0.0, 0.0))
    return (g + 1.0) / (g + b + 2.0)


def rate_arms(used, sign):
    """A rated generation: every choice the camera made in it shares the credit (sign > 0) or
    the blame (sign < 0), so a run with many choices does not punish each as hard as a run with
    one. `used`: [arm key]. -> arms rated."""
    used = [u for u in used or [] if isinstance(u, str) and u]
    if not used or not sign:
        return 0
    data = _read()
    arms = data.setdefault("arms", {})
    share = 1.0 / len(used)
    for u in used:
        g, b = arms.get(u, [0.0, 0.0])
        arms[u] = [round(g + (share if sign > 0 else 0.0), 4), round(b + (share if sign < 0 else 0.0), 4)]
    _save(data)
    return len(used)


def summary(limit=40):
    """What is remembered, for the person to read and prune: words (picked / replaced / merely
    recurring) and views (good / bad, with the per-trait rows folded under their view)."""
    data = _read()
    picks, rejects, seen = (data.get("picks") or {}), (data.get("rejects") or {}), (data.get("seen") or {})
    words = [{"word": w, "picks": int(picks.get(w, 0)), "rejects": int(rejects.get(w, 0)), "seen": int(seen.get(w, 0))}
             for w in {*picks, *rejects, *seen}]
    words.sort(key=lambda r: -(r["picks"] * 3 + r["rejects"] * 3 + r["seen"]))
    views = [{"view": k, "good": g, "bad": b} for k, (g, b) in view_stats().items() if "@" not in k]
    views.sort(key=lambda r: -(r["good"] + r["bad"]))
    return {"words": words[:limit], "more": max(0, len(words) - limit), "views": views,
            "prompts": int(data.get("prompts", 0)), "shots": int(data.get("shots", 0)),
            "kept": int(data.get("kept", 0))}


def forget(kind, name=None):
    """Drop one remembered word (`kind="word"`), one view and its per-trait rows (`"view"`), or
    everything (`"all"`). -> True if anything was removed."""
    data = _read()
    if kind == "all":
        existed = bool(data)
        data = {}
    elif kind == "word":
        existed = any(name in (data.get(k) or {}) for k in ("picks", "rejects", "seen"))
        for k in ("picks", "rejects", "seen"):
            (data.get(k) or {}).pop(name, None)
        existed = existed or (data.get("arms") or {}).pop(f"word:{name}", None) is not None
    elif kind == "view":
        views = data.get("views") or {}
        gone = [k for k in views if k == name or k.startswith(f"{name}@")]
        existed = bool(gone)
        for k in gone:
            views.pop(k)
    else:
        raise ValueError(f"unknown kind {kind!r}")
    _save(data)
    return existed


def effective_chance(chance):
    """The configured chance of a move, pulled toward how often the user actually keeps one
    (of the shots they reviewed) once enough were reviewed."""
    data = _read()
    total, kept = int(data.get("shots", 0)), int(data.get("kept", 0))
    if total >= MIN_SHOTS_FOR_RATE:
        chance = (chance + kept / total) / 2.0
    return max(0.05, min(1.0, chance * _lean("move:yes", "move:no")))


def _lean(yes, no):
    """Factor (0.5-2) that tilts an odds toward the arm the ratings liked more, 1 when unrated."""
    stats = arm_stats()
    ry, rn = arm_rate(yes, stats), arm_rate(no, stats)
    return 2.0 * ry / (ry + rn)


def split_chance(chance):
    """The chance of cutting a shot in two, tilted by whether split runs were liked."""
    return max(0.0, min(1.0, chance * _lean("split:yes", "split:no")))


def _save(data):
    path = _path()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f)
    os.replace(tmp, path)


def observe(prompt_hash, lemmas):
    """Count each of `lemmas` once for this prompt. -> prompts seen so far (0 = nothing kept)."""
    if not lemmas:
        return 0
    data = _read()
    hashes = data.get("hashes") or []
    if prompt_hash in hashes:
        return int(data.get("prompts", 0))
    seen = data.get("seen") or {}
    for lemma in set(lemmas):
        seen[lemma] = seen.get(lemma, 0) + 1
    data.update(seen=seen, prompts=int(data.get("prompts", 0)) + 1,
                hashes=(hashes + [prompt_hash])[-MAX_HASHES:])
    _save(data)
    return data["prompts"]