1jehuang--jcode/ios/TestHarness/reward/interaction/engine.py

"""engine.py - price the user-behavior graph into interaction-cost metrics.

Given the weighted ActionGraph (user_model) and target geometry (ui_map), this:

  1. Normalizes each state's out-edge weights into transition probabilities,
     forming a Markov chain over UI states.
  2. Solves for the stationary distribution pi (how often a user occupies each
     state over a long session) via power iteration.
  3. Prices every action in SECONDS with the KLM/TLM + Fitts cost_model, tracking
     the previous target so repeat-tap Fitts movement is charged correctly.
  4. Produces:
       - expected_action_cost_s: the expected seconds per action, weighting each
         action by pi[state] * P(action | state). This is the headline number:
         the average cost of a thing the user does, in real seconds.
       - task_times: completion time per canonical Task (sum of its action costs),
         and a frequency-weighted mean task time.
       - per_state and per_action detail for debugging / optimization targeting.

Lower seconds = better. To feed the reward (higher=better), call
reward_score(), which maps expected cost through a calibrated curve to 0..100.
"""

from __future__ import annotations

from dataclasses import dataclass, field, replace

from reward.interaction.cost_model import action_cost
from reward.interaction.model import Action, ActionGraph, Operators, UITarget
from reward.interaction.user_model import build_user_model


def _targets_for_action(a: Action, targets: dict[str, UITarget]) -> dict[str, UITarget]:
    """Return targets with the action's OWN target forced present.

    Some controls are context-conditional (e.g. the composer 'stop' button only
    appears while a turn is running). The ui_map marks those exists=False, but
    when the user performs the action that uses them, they are by definition on
    screen. Forcing just this action's target present avoids a spurious
    'unreachable' penalty while still flagging genuinely missing controls.
    """
    if not a.target_id or a.target_id not in targets:
        return targets
    t = targets[a.target_id]
    if t.exists:
        return targets
    patched = dict(targets)
    patched[a.target_id] = replace(t, exists=True)
    return patched



@dataclass
class EngineResult:
    expected_action_cost_s: float
    mean_task_time_s: float
    stationary: dict[str, float]
    action_costs_s: dict[str, float]
    action_probability: dict[str, float]   # pi[src] * P(action|src)
    task_times_s: dict[str, float]
    meta: dict = field(default_factory=dict)


def _stationary_distribution(graph: ActionGraph, iters: int = 500, tol: float = 1e-9) -> dict[str, float]:
    """Power-iterate the state-transition matrix to its stationary distribution.

    Transition prob from s = sum over out-edges to dst of normalized weight.
    Self-loops (chat->chat actions like send/scroll) keep mass in that state,
    correctly making 'chat' dominant since that is where most actions happen.
    """
    states = list(graph.states)
    idx = {s: i for i, s in enumerate(states)}
    n = len(states)

    # Build row-stochastic matrix P[s][dst].
    P = [[0.0] * n for _ in range(n)]
    for s in states:
        edges = graph.out_edges(s)
        total = sum(max(0.0, a.weight) for a in edges)
        if total <= 0:
            P[idx[s]][idx[s]] = 1.0   # absorbing if no edges
            continue
        for a in edges:
            P[idx[s]][idx[a.dst]] += max(0.0, a.weight) / total

    # Power iteration from uniform.
    pi = [1.0 / n] * n
    for _ in range(iters):
        nxt = [0.0] * n
        for i in range(n):
            pij = P[i]
            pii = pi[i]
            if pii == 0.0:
                continue
            for j in range(n):
                nxt[j] += pii * pij[j]
        # normalize + check convergence
        s = sum(nxt) or 1.0
        nxt = [v / s for v in nxt]
        if max(abs(nxt[k] - pi[k]) for k in range(n)) < tol:
            pi = nxt
            break
        pi = nxt
    return {states[i]: pi[i] for i in range(n)}


def run_engine(
    *,
    days: int = 7,
    log_dir: str = "~/.jcode/logs",
    source_root: str | None = None,
    ops: Operators = Operators(),
) -> EngineResult:
    graph, targets, meta = build_user_model(days=days, log_dir=log_dir, source_root=source_root)

    pi = _stationary_distribution(graph)

    # Price each action. Track previous target per source state so repeat taps
    # in the same state pay a realistic (small) movement cost.
    action_costs: dict[str, float] = {}
    prev_by_state: dict[str, str | None] = {s: None for s in graph.states}
    for a in graph.actions.values():
        eff = _targets_for_action(a, targets)
        bd = action_cost(a, eff, prev_target_id=prev_by_state.get(a.src), ops=ops)
        action_costs[a.id] = bd.seconds
        if a.target_id:
            prev_by_state[a.src] = a.target_id

    # Probability of each action = pi[src] * P(action | src).
    action_prob: dict[str, float] = {}
    for s in graph.states:
        edges = graph.out_edges(s)
        total = sum(max(0.0, e.weight) for e in edges) or 1.0
        for e in edges:
            action_prob[e.id] = pi.get(s, 0.0) * (max(0.0, e.weight) / total)

    # Expected seconds per action (the headline metric).
    psum = sum(action_prob.values()) or 1.0
    expected = sum(action_costs[aid] * p for aid, p in action_prob.items()) / psum

    # Task completion times (sum of action costs in the task's sequence).
    task_times: dict[str, float] = {}
    freq_weighted_num = 0.0
    freq_weighted_den = 0.0
    for t in graph.tasks:
        secs = 0.0
        prev: str | None = None
        for aid in t.action_ids:
            a = graph.actions[aid]
            eff = _targets_for_action(a, targets)
            bd = action_cost(a, eff, prev_target_id=prev, ops=ops)
            secs += bd.seconds
            if a.target_id:
                prev = a.target_id
        task_times[t.id] = round(secs, 3)
        freq_weighted_num += secs * max(0.0, t.frequency)
        freq_weighted_den += max(0.0, t.frequency)
    mean_task = (freq_weighted_num / freq_weighted_den) if freq_weighted_den else 0.0

    return EngineResult(
        expected_action_cost_s=round(expected, 4),
        mean_task_time_s=round(mean_task, 4),
        stationary={k: round(v, 4) for k, v in pi.items()},
        action_costs_s={k: round(v, 4) for k, v in action_costs.items()},
        action_probability={k: round(v, 5) for k, v in action_prob.items()},
        task_times_s=task_times,
        meta=meta,
    )


# Calibration for mapping seconds -> 0..100 reward (higher = cheaper to use).
# Anchors chosen from the model's own range: an expected per-action cost of
# ~1.5s (mostly cheap taps/reads) is excellent (->100); ~6s (deep, slow flows)
# is poor (->0). Linear in between, clamped.
_GOOD_S = 1.5
_BAD_S = 6.0


def reward_score(result: EngineResult | None = None, **kwargs) -> float:
    """Map the engine's expected per-action cost to a 0..100 reward."""
    if result is None:
        result = run_engine(**kwargs)
    s = result.expected_action_cost_s
    if s <= _GOOD_S:
        return 100.0
    if s >= _BAD_S:
        return 0.0
    return round(100.0 * (_BAD_S - s) / (_BAD_S - _GOOD_S), 2)


if __name__ == "__main__":
    import json

    r = run_engine()
    print("usage source:", r.meta.get("usage_source"), "lines:", r.meta.get("lines_scanned"))
    print(f"expected per-action cost: {r.expected_action_cost_s:.3f} s")
    print(f"freq-weighted mean task time: {r.mean_task_time_s:.3f} s")
    print(f"reward score: {reward_score(r):.1f}/100")
    print("stationary distribution:", json.dumps(r.stationary))
    print("task times (s):", json.dumps(r.task_times_s))
    print("action costs (s):")
    for aid, c in sorted(r.action_costs_s.items(), key=lambda kv: -kv[1]):
        print(f"  {aid:16} {c:6.3f}  (p={r.action_probability.get(aid,0):.4f})")