app/schedule_math.py

"""Pure schedule math: duration shrinkage, slip, completion rates, EV score.

Server-owned — LLM must never invent these statistics.

P1 scope lock (do not expand):
- EV stays exactly:
  EV = P(done)*U + γ*I + η*Fun − λt*d_hat − λm*cost − λf*FSE
- No PERMA, no ωE_i, no δ_i, no H=S+C+V in code.
- Hedonic δ is not a v1 term; learn later via would_repeat + falling fun.
- Seligman = operators on a trigger day: raise λ_f and η; only P0 overrides.
- Escape (porn/court/rerun binge) ≠ explore/restore_fun.
- Family contact is often FSE load, not a pillar to maximize.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Any

# Default planned minutes by kind when no samples exist yet.
PRIOR_DEFAULT_MIN: dict[str, float] = {
    "earn_ship": 45,
    "admin_spain": 30,
    "body_care": 20,
    "move_out": 15,
    "boundary": 10,
    "food_out": 45,
    "stabilize": 15,
    "explore": 25,
    "restore_fun": 30,
    "sleep_window": 45,
    "other": 30,
}

# Escape / FSE tags — never treat as fun or explore intent.
ESCAPE_FSE_TAGS = frozenset(
    {"urge", "corn", "court", "rerun", "daydream", "bully", "shame", "family"}
)


def is_strong_feedback(fb: dict[str, Any]) -> bool:
    """Strong iff done|partial with actual_min, quality, and fun present."""

    did = fb.get("did")
    return (
        did in ("done", "partial")
        and fb.get("actual_min") is not None
        and fb.get("quality") is not None
        and fb.get("fun") is not None
    )


def d_hat(
    mean_actual: float | None,
    n: int,
    prior_default: float,
    *,
    m: float = 3.0,
) -> float:
    """Shrinkage duration estimate: (n/(n+m))*mean + (m/(n+m))*prior."""

    if n <= 0 or mean_actual is None:
        return float(prior_default)
    weight = n / (n + m)
    return weight * float(mean_actual) + (1.0 - weight) * float(prior_default)


def p_done_rate(n_done: int, n: int, *, alpha: float = 1.0, beta: float = 1.0) -> float:
    """Smoothed completion rate (N_done + α) / (N + α + β)."""

    return (n_done + alpha) / (n + alpha + beta) if (n + alpha + beta) else 0.5


def mean_or_none(values: list[float]) -> float | None:
    return sum(values) / len(values) if values else None


def recompute_priors(
    feedback: list[dict[str, Any]],
    blocks_by_id: dict[str, dict[str, Any]],
    *,
    shrink_k: float = 3.0,
) -> dict[str, dict[str, Any]]:
    """Rebuild per-kind priors from strong feedback samples.

    EV ranking (documented for schedule generators / OR prompts):
      EV ≈ P_done * U + γ * I + η * Fun_bar - λt * d_hat - λm * cost - λf * fse_load
    High I when n_k is low or variance is high (explore kinds).
    """

    by_kind: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for fb in feedback:
        if not (fb.get("strong") or is_strong_feedback(fb)):
            continue
        block = blocks_by_id.get(str(fb.get("block_id") or ""))
        kind = (block or {}).get("kind") or fb.get("kind") or "other"
        by_kind[str(kind)].append(fb)

    out: dict[str, dict[str, Any]] = {}
    for kind, samples in by_kind.items():
        actuals = [float(s["actual_min"]) for s in samples if s.get("actual_min") is not None]
        qualities = [float(s["quality"]) for s in samples if s.get("quality") is not None]
        funs = [float(s["fun"]) for s in samples if s.get("fun") is not None]
        energies = [
            float(s["energy_after"])
            for s in samples
            if s.get("energy_after") is not None
        ]
        slips: list[float] = []
        for s in samples:
            block = blocks_by_id.get(str(s.get("block_id") or ""))
            planned = (block or {}).get("planned_min")
            if planned is not None and s.get("actual_min") is not None:
                slips.append(float(s["actual_min"]) - float(planned))
        n = len(samples)
        n_done = sum(1 for s in samples if s.get("did") == "done")
        mean_actual = mean_or_none(actuals)
        prior = PRIOR_DEFAULT_MIN.get(kind, 30.0)
        repeats = [
            1.0 if s.get("would_repeat") == "yes" else 0.0
            for s in samples
            if s.get("would_repeat") in ("yes", "no", "maybe")
        ]
        out[kind] = {
            "kind": kind,
            "n": n,
            "mean_actual": mean_actual,
            "d_hat": d_hat(mean_actual, n, prior, m=shrink_k),
            "mean_quality": mean_or_none(qualities),
            "mean_fun": mean_or_none(funs),
            "mean_energy": mean_or_none(energies),
            "p_done": p_done_rate(n_done, n),
            "mean_slip": mean_or_none(slips),
            "repeat_score": mean_or_none(repeats) or 0.0,
        }
    # Ensure defaults exist for known kinds with no samples.
    for kind, prior in PRIOR_DEFAULT_MIN.items():
        if kind not in out:
            out[kind] = {
                "kind": kind,
                "n": 0,
                "mean_actual": None,
                "d_hat": float(prior),
                "mean_quality": None,
                "mean_fun": None,
                "mean_energy": None,
                "p_done": 0.5,
                "mean_slip": None,
                "repeat_score": 0.0,
            }
    return out


def minutes_between(start: str, end: str) -> int:
    """Compute end-start in minutes for HH:MM strings."""

    sh, sm = map(int, start.split(":"))
    eh, em = map(int, end.split(":"))
    return (eh * 60 + em) - (sh * 60 + sm)


def parse_hhmm(value: str) -> int:
    h, m = map(int, value.split(":"))
    return h * 60 + m


def overlaps(a_start: str, a_end: str, b_start: str, b_end: str) -> bool:
    a0, a1 = parse_hhmm(a_start), parse_hhmm(a_end)
    b0, b1 = parse_hhmm(b_start), parse_hhmm(b_end)
    return a0 < b1 and b0 < a1


def validate_blocks(
    blocks: list[dict[str, Any]],
    *,
    max_blocks: int = 7,
    previous_p0: list[dict[str, Any]] | None = None,
    allow_p0_move: bool = False,
    must_include_explore_or_restore: bool = False,
    capacity_hint: float | None = None,
    hard_explore: bool = False,
) -> tuple[list[str], list[str]]:
    """Return (errors, warnings). Empty errors means valid."""

    errors: list[str] = []
    warnings: list[str] = []
    if len(blocks) > max_blocks:
        errors.append(f"max_blocks exceeded ({len(blocks)} > {max_blocks})")
    for block in blocks:
        start = str(block.get("start") or "")
        end = str(block.get("end") or "")
        try:
            span = minutes_between(start, end)
        except Exception:  # noqa: BLE001
            errors.append(f"invalid time on block {block.get('id')}")
            continue
        if span <= 0:
            errors.append(f"end must be after start for {block.get('id')}")
        planned = int(block.get("planned_min") or 0)
        if planned and abs(planned - span) > 1:
            errors.append(
                f"planned_min mismatch for {block.get('id')}: {planned} vs {span}"
            )
    for i, a in enumerate(blocks):
        for b in blocks[i + 1 :]:
            if overlaps(
                str(a.get("start")),
                str(a.get("end")),
                str(b.get("start")),
                str(b.get("end")),
            ):
                errors.append(
                    f"overlap between {a.get('id')} and {b.get('id')}"
                )
    if previous_p0 and not allow_p0_move:
        prev = {
            str(b.get("id")): b
            for b in previous_p0
            if b.get("priority") == "P0" or b.get("locked")
        }
        for block in blocks:
            bid = str(block.get("id") or "")
            if bid in prev:
                old = prev[bid]
                if old.get("start") != block.get("start") or old.get("end") != block.get(
                    "end"
                ):
                    errors.append(f"P0/locked block {bid} cannot be moved")
    if (
        must_include_explore_or_restore
        and blocks
        and capacity_hint is not None
        and capacity_hint >= 0.4
    ):
        has_explore = any(
            b.get("intent") in ("explore", "restore_fun") for b in blocks
        )
        if not has_explore:
            msg = "all-grind plan: add explore or restore_fun when capacity allows"
            if hard_explore:
                errors.append(msg)
            else:
                warnings.append(msg)
    return errors, warnings


def ev_score(
    *,
    p_done: float,
    utility: float,
    information: float = 0.0,
    fun: float = 0.0,
    d_hat_min: float = 30.0,
    cost: float = 0.0,
    fse: float = 0.0,
    gamma: float = 0.25,
    eta: float = 0.35,
    lambda_t: float = 0.01,
    lambda_m: float = 0.15,
    lambda_f: float = 0.4,
    raise_eta: bool = False,
    raise_lambda_f: bool = False,
) -> float:
    """Locked EV formula (operators may raise η / λ_f; no δ / H / PERMA terms)."""

    eta_eff = eta * (1.35 if raise_eta else 1.0)
    lf_eff = lambda_f * (1.4 if raise_lambda_f else 1.0)
    return (
        p_done * utility
        + gamma * information
        + eta_eff * fun
        - lambda_t * d_hat_min
        - lambda_m * cost
        - lf_eff * fse
    )


def capacity_hint(
    *,
    risk_1h_score: float,
    triggers_yesterday: list[str],
    yesterday_trigger: bool | None = None,
    last_hour_high: bool | None = None,
) -> float:
    """Soft capacity 0..1. Trigger day / last-hour FSE cut C_today; P0 still allowed.

    On trigger days, callers should also raise λ_f and η in ranking/reschedule
    prompts (operators only — EV formula unchanged).
    """

    y_trig = (
        bool(yesterday_trigger)
        if yesterday_trigger is not None
        else len(triggers_yesterday) > 0
    )
    hour_high = (
        bool(last_hour_high)
        if last_hour_high is not None
        else risk_1h_score >= 1.0
    )
    penalty = min(
        0.65,
        0.15 * risk_1h_score
        + 0.08 * len(triggers_yesterday)
        + (0.12 if y_trig else 0.0)
        + (0.15 if hour_high else 0.0),
    )
    return max(0.25, 1.0 - penalty)


def trigger_operators(
    *,
    risk_1h_tags: list[str],
    triggers_yesterday: list[str],
) -> dict[str, Any]:
    """Seligman-as-operators payload for reschedule / agent context (not new math)."""

    escape_hit = bool(ESCAPE_FSE_TAGS & set(risk_1h_tags + triggers_yesterday))
    trigger_day = len(triggers_yesterday) > 0 or escape_hit
    return {
        "trigger_day": trigger_day,
        "raise_lambda_f": trigger_day,
        "raise_eta": trigger_day,
        "p0_only_overrides": True,
        "escape_is_not_explore": True,
        "fun_must_yield_data_or_skill": True,
        "no_family_relationship_optimization": True,
        "ev_formula_locked": True,
    }


def plan_health(
    blocks: list[dict[str, Any]],
    feedback: list[dict[str, Any]],
) -> dict[str, Any]:
    """Health ≈ (strong_fb/planned) * p0_done_rate * explore_flag."""

    planned = [b for b in blocks if b.get("status") != "cancelled"]
    planned_count = max(1, len(planned))
    fb_by_block = {str(f.get("block_id")): f for f in feedback}
    strong = sum(
        1
        for b in planned
        if is_strong_feedback(fb_by_block.get(str(b.get("id")), {}))
        or fb_by_block.get(str(b.get("id")), {}).get("strong")
    )
    p0 = [b for b in planned if b.get("priority") == "P0"]
    p0_done = sum(
        1
        for b in p0
        if fb_by_block.get(str(b.get("id")), {}).get("did") == "done"
        or b.get("status") == "done"
    )
    p0_rate = (p0_done / len(p0)) if p0 else 1.0
    explore_done = any(
        b.get("intent") in ("explore", "restore_fun")
        and (
            b.get("status") in ("done", "partial")
            or fb_by_block.get(str(b.get("id")), {}).get("did") in ("done", "partial")
        )
        for b in planned
    )
    score = (strong / planned_count) * p0_rate * (1.0 if explore_done else 0.0)
    return {
        "score": round(score, 3),
        "strong_feedback_count": strong,
        "planned_count": len(planned),
        "p0_done_rate": round(p0_rate, 3),
        "explore_or_restore_done": explore_done,
    }


def priors_markdown(priors: dict[str, dict[str, Any]]) -> str:
    """Compact SERVER_PRIORS table for LLM prompts."""

    lines = [
        "SERVER_PRIORS (do not invent numbers)",
        "kind | n | d_hat | p_done | fun_bar | repeat",
    ]
    for kind in sorted(priors.keys()):
        row = priors[kind]
        fun = row.get("mean_fun")
        fun_s = f"{fun:.2f}" if isinstance(fun, (int, float)) else "n/a"
        lines.append(
            f"{kind} | {row.get('n', 0)} | {row.get('d_hat', 0):.0f} | "
            f"{row.get('p_done', 0):.2f} | {fun_s} | {row.get('repeat_score', 0):.2f}"
        )
    return "\n".join(lines)