simulation.py

"""SIMULATION-mode backend for the Stress Test the Octopus demo.

This module owns the demo's *state* and the *real* monitoring pipeline.
FI values are simulated (baseline 0.07 %, pushed up by noise injection),
but they are fed through the genuine
:class:`octopus.monitoring.alerts.FIAlertSystem` and
:class:`octopus.monitoring.regulator.SelfRegulator` — so the alerts and
the 5-stage lifecycle on screen are produced by the same code the
platform ships, not faked.

Decomposition + routing also use the real
:class:`octopus.router.task_router.TaskRouter` (rule-based path), and it
is made *topology-aware*: an arm the operator disabled, or one the
regulator has isolated, is treated as unavailable, so generation falls
back exactly as the platform would.
"""

from __future__ import annotations

import logging
import os
import random
from dataclasses import dataclass, field
from datetime import datetime

from octopus.monitoring import (
    FIAlertSystem,
    SelfRegulator,
    SelfRegulatorConfig,
)
from octopus.monitoring.alerts import (
    KIND_CRITICAL,
    KIND_DEGRADED,
    KIND_HEALTHY,
    KIND_RECOVERED,
)
from octopus.monitoring.regulator import (
    ArmStage,
    STRATEGY_ARM_ISOLATION,
    STRATEGY_FALLBACK_MODE,
    STRATEGY_ROUTE_SHIFT,
)
from octopus.router.task_router import TaskRouter

import demo_data

logger = logging.getLogger(__name__)

# The four arms, in display order. code_review has no disable button (it
# isn't exercised by the example prompts) but is still monitored.
ARMS = ["code_generation", "testing", "code_review", "cicd"]
DISABLEABLE = ["code_generation", "testing", "cicd"]

# Friendly labels for the UI.
ARM_LABELS = {
    "code_generation": "code_gen",
    "testing": "testing",
    "code_review": "code_review",
    "cicd": "cicd",
}

BASELINE_FI = 0.07  # percent — healthy structural baseline

# FI the targeted arm climbs to after the Nth noise injection. Tuned so a
# short demo walks the whole lifecycle: 1st → DEGRADED, 2nd → ISOLATED,
# 4th → FALLBACK. (Thresholds below: warn/degrade 1 %, crit/isolate 5 %,
# fallback 20 %.)
_NOISE_LEVELS = [2.5, 6.0, 13.0, 26.0, 42.0, 55.0]

# Demo thresholds, aligned with the gauge bands (green <1, yellow 1-5,
# red >5) so what the operator sees matches what the regulator does.
_ALERT_WARNING = 1.0
_ALERT_CRITICAL = 5.0
_REG_CONFIG = SelfRegulatorConfig(
    degrade_activate_pct=1.0, degrade_recover_pct=0.5,
    isolate_activate_pct=5.0, isolate_recover_pct=3.0,
    fallback_activate_pct=20.0, fallback_recover_pct=15.0,
)


def detect_mode(force_live: bool = False) -> str:
    """Return ``"LIVE"`` or ``"SIMULATION"``.

    Auto-detects **LIVE** when trained checkpoints are present *and* a CUDA GPU
    is available; otherwise **SIMULATION**. ``force_live`` (or ``OCTOPUS_LIVE=1``)
    requests LIVE whenever checkpoints exist — so a GPU pod can be forced on even
    if the auto-check is conservative. Without checkpoints it always stays
    SIMULATION (the CPU/no-weights case).

    Checkpoint availability is delegated to :mod:`live_backend` (which knows the
    configured local dir / hub repo); the check is filesystem-only and never
    loads a model.
    """
    try:
        import live_backend
        has_ckpt = live_backend.checkpoints_available()
    except Exception:
        has_ckpt = False
    try:
        import torch
        has_gpu = torch.cuda.is_available()
    except Exception:
        has_gpu = False

    if not has_ckpt:
        return "SIMULATION"
    if force_live or os.environ.get("OCTOPUS_LIVE") == "1":
        return "LIVE"
    return "LIVE" if has_gpu else "SIMULATION"


def _live_backend():
    """The module-global LiveBackend singleton (lazy import; never deep-copied)."""
    import live_backend
    return live_backend.get_backend()


MODE_BANNER = {
    "SIMULATION": "Demo mode — using pre-computed outputs (no GPU)",
    "LIVE": "Live inference — Mistral 7B + 4 trained arms",
}


@dataclass
class FeedLine:
    """One line in the alert feed."""

    time: str
    tag: str      # e.g. ARM_DEGRADED, ROUTE_SHIFT, ARM_ISOLATED, RECOVERY
    text: str
    level: str    # info | warning | critical | recovery (drives colour)


@dataclass
class DemoState:
    """All per-session demo state. Recreated on Restore.

    Holds the live monitoring objects so their hysteresis/lifecycle state
    persists across button clicks within a session.
    """

    mode: str = "SIMULATION"
    fi: dict = field(default_factory=lambda: {a: BASELINE_FI for a in ARMS})
    disabled: set = field(default_factory=set)
    noise_count: dict = field(default_factory=lambda: {a: 0 for a in ARMS})
    feed: list = field(default_factory=list)
    attacks: int = 0
    alerts: FIAlertSystem = field(
        default_factory=lambda: FIAlertSystem(_ALERT_WARNING, _ALERT_CRITICAL)
    )
    regulator: SelfRegulator = field(
        default_factory=lambda: SelfRegulator(_REG_CONFIG)
    )
    _rng: random.Random = field(default_factory=lambda: random.Random(7))

    def __post_init__(self) -> None:
        # LIVE mode delegates model ops to the module-global live backend. The
        # backend is NOT stored on the state (it holds GB of weights, and
        # gr.State deep-copies the state per session) — we reach it via
        # live_backend.get_backend() on demand instead.
        self._live = self.mode == "LIVE"
        # Register every arm with the monitoring stack at baseline so the status
        # cards have a stage from the first render. In LIVE mode the *real* FI is
        # pulled from the backend on the first state-changing action, so we don't
        # load the 7B model at app-startup / first render.
        self._observe(log=False)

    def _apply_live(self, op: str, arm: "str | None" = None) -> bool:
        """Drive a LIVE-mode model operation and refresh ``self.fi`` from the
        real per-arm Fragility Index.

        Returns ``False`` (so the caller falls back to the SIMULATION behaviour)
        when the backend is unavailable or errors — the demo never hard-crashes.
        """
        try:
            backend = _live_backend()
            if op == "inject_noise":
                backend.inject_noise(arm, intensity=1.0)
            elif op == "restore_all":
                backend.restore_all()
            elif op == "disable":
                backend.arm_disable(arm)
            elif op == "enable":
                backend.arm_enable(arm)
            if op in ("inject_noise", "restore_all"):
                fi_map = backend.get_fi()
                for name, value in (fi_map or {}).items():
                    if name in self.fi:
                        self.fi[name] = round(float(value), 4)
            return True
        except Exception:
            logger.exception("LIVE backend op %r failed; using simulation fallback", op)
            return False

    # ------------------------------------------------------------------
    # Status helpers (read by the UI)
    # ------------------------------------------------------------------

    def arm_status(self, arm: str) -> str:
        """UI status label for an arm."""
        if arm in self.disabled:
            return "Disabled"
        stage = self.regulator.arm_stages.get(arm, ArmStage.HEALTHY)
        return {
            ArmStage.HEALTHY: "Active",
            ArmStage.DEGRADED: "Degraded",
            ArmStage.ISOLATED: "Isolated",
            ArmStage.FALLBACK: "Fallback",
            ArmStage.RECOVERING: "Recovering",
        }[stage]

    def is_available(self, arm: str) -> bool:
        """An arm is routable if not manually disabled and not pulled by
        the regulator (ISOLATED / FALLBACK)."""
        if arm in self.disabled:
            return False
        return arm not in self.regulator.isolated_arms

    def cross_domain_impact(self) -> float:
        """Max FI rise (percentage points) on any arm that was *not*
        attacked — the headline "did the damage spread?" number."""
        attacked = {a for a, c in self.noise_count.items() if c > 0}
        others = [a for a in ARMS if a not in attacked]
        if not others:
            return 0.0
        return max(0.0, max(self.fi[a] - BASELINE_FI for a in others))

    def mission_summary(self) -> str:
        impact = self.cross_domain_impact()
        # Every attack is "survived" while the blast stayed contained
        # (cross-domain impact ~0) — the regulator isolated the arm.
        survived = self.attacks if impact < 0.5 else max(0, self.attacks - 1)
        return (
            f"Attacks: {self.attacks} | Survived: {survived} | "
            f"Cross-domain impact: {impact:.2f}%"
        )

    # ------------------------------------------------------------------
    # Stress-test actions
    # ------------------------------------------------------------------

    def toggle_disable(self, arm: str) -> None:
        if arm in self.disabled:
            self.disabled.discard(arm)
            if self._live:
                self._apply_live("enable", arm)
            self._log("MANUAL", f"{ARM_LABELS[arm]} re-enabled by operator", "info")
        else:
            self.disabled.add(arm)
            if self._live:
                self._apply_live("disable", arm)
            self._log("MANUAL", f"{ARM_LABELS[arm]} disabled by operator", "warning")

    def inject_noise(self, arm: str = "code_generation") -> None:
        """Add Gaussian noise to one arm, pushing its FI up.

        SIMULATION: FI follows a scripted escalation. LIVE: real Gaussian noise
        is added to the arm's LoRA weights and the new FI is measured on the
        model — then run through the same monitoring/regulator code.
        """
        self.noise_count[arm] += 1
        if not (self._live and self._apply_live("inject_noise", arm)):
            idx = self.noise_count[arm] - 1
            if idx < len(_NOISE_LEVELS):
                base = _NOISE_LEVELS[idx]
            else:
                base = _NOISE_LEVELS[-1] + 12.0 * (idx - len(_NOISE_LEVELS) + 1)
            jitter = self._rng.gauss(0.0, 0.4)
            self.fi[arm] = max(0.0, round(base + jitter, 2))
        self.attacks += 1
        self._log(
            "INJECT",
            f"Gaussian noise injected into {ARM_LABELS[arm]} "
            f"(FI now {self.fi[arm]:.2f}%)",
            "critical",
        )
        self._observe()

    def restore_all(self) -> None:
        """Reset everything to a healthy baseline, walking arms back down
        through the lifecycle so RECOVERY events show in the feed."""
        self.disabled = set()
        self.noise_count = {a: 0 for a in ARMS}
        self.fi = {a: BASELINE_FI for a in ARMS}
        if self._live:
            # Reloads every arm's weights from snapshot + refreshes the real FI.
            self._apply_live("restore_all")
        # Two observations: 1st drives elevated arms → RECOVERING, 2nd
        # settles them → HEALTHY (matches the regulator's de-escalation).
        self._observe()
        self._observe()
        self.attacks = 0
        self._log("RESTORE", "all arms reset to healthy baseline", "recovery")

    # ------------------------------------------------------------------
    # Monitoring pipeline (the real alert + regulator code)
    # ------------------------------------------------------------------

    def _observe(self, log: bool = True) -> None:
        """Feed the current FI map through the alert system and regulator,
        turning their events into feed lines."""
        reading = dict(self.fi)
        alert_events = self.alerts.on_fi_update(reading)
        reg_events = self.regulator.on_fi_update(reading)
        if not log:
            return
        for a in alert_events:
            tag, text, level = _format_alert(a)
            self._log(tag, text, level)
        for e in reg_events:
            tag, text, level = _format_regulation(e)
            self._log(tag, text, level)

    def _log(self, tag: str, text: str, level: str) -> None:
        self.feed.append(
            FeedLine(time=datetime.now().strftime("%H:%M:%S"),
                     tag=tag, text=text, level=level)
        )


def _format_alert(alert) -> tuple[str, str, str]:
    arm = ARM_LABELS.get(alert.arm_name, alert.arm_name or "system")
    fi = alert.fi_pct
    if alert.kind == KIND_DEGRADED:
        return "ARM_DEGRADED", f"{arm} — FI rising to {fi:.2f}%", "warning"
    if alert.kind == KIND_CRITICAL:
        return "ARM_CRITICAL", f"{arm} — FI critical at {fi:.2f}%", "critical"
    if alert.kind == KIND_RECOVERED:
        return "RECOVERED", f"{arm} — FI easing to {fi:.2f}%", "recovery"
    if alert.kind == KIND_HEALTHY:
        return "HEALTHY", f"{arm} — FI back to {fi:.2f}%", "recovery"
    return alert.kind, f"{arm} — FI {fi:.2f}%", "info"


def _format_regulation(event) -> tuple[str, str, str]:
    arm = ARM_LABELS.get(event.arm_name, event.arm_name)
    fi = event.fi_pct
    if event.strategy == STRATEGY_ROUTE_SHIFT:
        return "ROUTE_SHIFT", f"{arm} traffic reduced", "warning"
    if event.strategy == STRATEGY_ARM_ISOLATION:
        return "ARM_ISOLATED", f"{arm} removed from routing", "critical"
    if event.strategy == STRATEGY_FALLBACK_MODE:
        return "FALLBACK", f"{arm} capability served brain-only", "critical"
    # RECOVERY strategy — distinguish "easing back" from "fully restored".
    if event.to_stage == ArmStage.HEALTHY:
        return "RECOVERY", f"{arm} restored — FI back to {fi:.2f}%", "recovery"
    return "RECOVERY", f"{arm} recovering — FI {fi:.2f}%", "recovery"


# ----------------------------------------------------------------------
# Generation pipeline (decompose -> dispatch -> attach pre-computed code)
# ----------------------------------------------------------------------

@dataclass
class GeneratedSubtask:
    """One produced subtask, ready to render in the output panel."""

    arm: str               # arm that handled it (may differ from target)
    target_arm: str        # arm the task type maps to
    confidence: float
    title: str
    filename: str
    language: str
    code: str
    fallback_note: str = ""  # set when the target arm was unavailable


def plan_generation(state: DemoState, instruction: str) -> list[GeneratedSubtask]:
    """Decompose + route an instruction and attach pre-computed outputs.

    Uses the curated answer for a recognised example prompt; otherwise the
    real rule-based router decides the subtasks. Either way, routing is
    topology-aware: a disabled/isolated target arm triggers the platform's
    fallback (reroute to code_generation, else brain-only).

    In LIVE mode this runs the real SoloWorkflow (decompose → route → arm
    inference → assemble) via the live backend, returning the same
    ``list[GeneratedSubtask]`` shape; it falls back to the SIMULATION path if
    live generation errors.
    """
    if getattr(state, "_live", False):
        try:
            backend = _live_backend()
            backend.sync_availability({a: state.is_available(a) for a in ARMS})
            subs = backend.generate(instruction)
            if subs:
                return subs
        except Exception:
            logger.exception("LIVE generation failed; using simulation fallback")

    curated = demo_data.curated_subtasks(instruction)
    router = TaskRouter()
    results: list[GeneratedSubtask] = []

    if curated is not None:
        for card in curated:
            target = card["arm"]
            arm, note = _resolve_arm(state, target)
            results.append(GeneratedSubtask(
                arm=arm or "brain", target_arm=target,
                confidence=card["confidence"], title=card["title"],
                filename=card["filename"], language=card["language"],
                code=card["code"], fallback_note=note,
            ))
        return results

    # Arbitrary instruction → real rule-based decomposition + dispatch.
    subtasks = router.decompose(instruction)
    for st in subtasks:
        router.dispatch(st, is_available=state.is_available)
        target = router.arm_for_task.get(st.task_type, "code_generation")
        arm = st.arm  # already resolved by dispatch (may be None = brain-only)
        note = ""
        if st.metadata.get("fallback"):
            fb = st.metadata["fallback"]
            to = ARM_LABELS.get(fb["to"], "brain-only") if fb["to"] else "brain-only"
            note = f"{ARM_LABELS.get(fb['from'], fb['from'])} unavailable → {to}"
        snippet = demo_data.generic_snippet(st.task_type.value)
        results.append(GeneratedSubtask(
            arm=arm or "brain", target_arm=target,
            confidence=round(st.confidence, 4),
            title=st.description[:80],
            filename=snippet["filename"], language=snippet["language"],
            code=snippet["code"], fallback_note=note,
        ))
    return results


def _resolve_arm(state: DemoState, target: str) -> tuple[str | None, str]:
    """Topology-aware arm resolution for a curated subtask.

    Returns ``(arm_used, fallback_note)``. Mirrors TaskRouter.dispatch:
    unavailable target → code_generation if *it* is available, else
    brain-only.
    """
    if state.is_available(target):
        return target, ""
    if target != "code_generation" and state.is_available("code_generation"):
        return (
            "code_generation",
            f"{ARM_LABELS[target]} unavailable → {ARM_LABELS['code_generation']}",
        )
    return None, f"{ARM_LABELS[target]} unavailable → brain-only fallback"