#!/usr/bin/env python3
"""
helix_cdc/runtime/policy_harness.py

WO-ECHO-POLICY-HARNESS-01 - Policy Harness for Runtime Optimization

Learns to optimize runtime knobs from telemetry. The harness observes metrics
(tok/s, cache hits, regen calls, latencies) and adjusts tunable knobs to
maximize throughput while maintaining correctness.

Key Principle: Reality is the judge. Receipts prove. Policy evolves only if
correctness is maintained.

Usage:
    from helix_cdc.runtime.policy_harness import PolicyHarness, PolicyState, PolicyAction

    # Load trained policy
    harness = PolicyHarness.from_checkpoint("receipts/policy/trained_v0.pt")

    # Extract state from telemetry
    state = harness.extract_state(telem)

    # Select action
    action = harness.select_action(state)

    # Apply action during inference
    with harness.apply_action(action):
        output, receipt, _ = stream_transformer_block_forward(...)

    # Record for training
    telem.record_policy_action(action.to_dict())
"""

from __future__ import annotations

import json
import os
from contextlib import contextmanager
from dataclasses import dataclass, asdict, field
from pathlib import Path
from typing import Any, Dict, Generator, List, Optional, Tuple, TYPE_CHECKING
import numpy as np

# Lazy torch import to avoid startup cost when not training
_torch = None
_nn = None


def _import_torch():
    """Lazy import torch."""
    global _torch, _nn
    if _torch is None:
        import torch
        import torch.nn as nn
        _torch = torch
        _nn = nn
    return _torch, _nn


if TYPE_CHECKING:
    from .telemetry import Telemetry


# =============================================================================
# Constants: Knob Value Mappings
# =============================================================================

# WO-ECHO-POLICY-GOLDEN-01: Reduced action space (4 wirable knobs only)
# Removed: sample_rate (hardcoded 0.1), cache_size (no param exists),
#          block_rows (burned into CDNA format at encoding time)

# Discrete value mappings for each knob
VERIFY_POLICY_VALUES = ["always", "sampled", "trust_cached"]
FIBPI_MODE_VALUES = ["round_robin", "se_bucket"]

# Head dimensions for multi-head output (reduced from 7 to 4 knobs)
HEAD_DIMS = {
    "verify_policy": 3,   # 0=always, 1=sampled, 2=trust_cached
    "cpp_kernel": 2,      # 0=off, 1=on
    "fused_matmul": 2,    # 0=off, 1=on
    "fibpi_mode": 2,      # 0=round_robin, 1=se_bucket
}

# Total action space size: 3 × 2 × 2 × 2 = 24 (was 2400)
TOTAL_ACTION_COMBINATIONS = 3 * 2 * 2 * 2  # = 24

# Deprecated: Kept for backwards compatibility when loading old checkpoints
_DEPRECATED_HEAD_DIMS = {
    "sample_rate": 5,
    "cache_size": 4,
    "block_rows": 5,
}


# =============================================================================
# PolicyState: 14-feature state vector
# =============================================================================

@dataclass
class PolicyState:
    """
    14-feature state vector extracted from telemetry snapshots.

    Features are organized into groups:
    - Performance (4): tok_per_sec, mean_step_ms, p95_step_ms, steps_in_window
    - Cache (3): cache_hit_rate, cache_misses, cache_hits
    - Regen (3): regen_hit_rate, regen_time_pct, regen_calls
    - Memory (2): rss_delta_mb, rss_mb
    - Se (2): se_mean, se_std
    """
    # Performance (4)
    tok_per_sec: float = 0.0
    mean_step_ms: float = 0.0
    p95_step_ms: float = 0.0
    steps_in_window: int = 0

    # Cache (3)
    cache_hit_rate: float = 0.0
    cache_misses: int = 0
    cache_hits: int = 0

    # Regen (3)
    regen_hit_rate: float = 1.0
    regen_time_pct: float = 0.0
    regen_calls: int = 0

    # Memory (2)
    rss_delta_mb: float = 0.0
    rss_mb: float = 0.0

    # Se (2)
    se_mean: float = 0.0
    se_std: float = 0.0

    def to_array(self) -> np.ndarray:
        """Convert to numpy array (14 features)."""
        return np.array([
            self.tok_per_sec,
            self.mean_step_ms,
            self.p95_step_ms,
            float(self.steps_in_window),
            self.cache_hit_rate,
            float(self.cache_misses),
            float(self.cache_hits),
            self.regen_hit_rate,
            self.regen_time_pct,
            float(self.regen_calls),
            self.rss_delta_mb,
            self.rss_mb,
            self.se_mean,
            self.se_std,
        ], dtype=np.float32)

    def to_tensor(self):
        """Convert to torch tensor (14 features)."""
        torch, _ = _import_torch()
        return torch.tensor(self.to_array(), dtype=torch.float32)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "PolicyState":
        """Create from dict (e.g., telemetry snapshot)."""
        return cls(
            tok_per_sec=float(d.get("tok_per_sec", 0.0)),
            mean_step_ms=float(d.get("mean_step_ms", 0.0)),
            p95_step_ms=float(d.get("p95_step_ms", 0.0)),
            steps_in_window=int(d.get("steps_in_window", 0)),
            cache_hit_rate=float(d.get("cache_hit_rate", 0.0)),
            cache_misses=int(d.get("cache_misses", 0)),
            cache_hits=int(d.get("cache_hits", 0)),
            regen_hit_rate=float(d.get("regen_hit_rate", 1.0)),
            regen_time_pct=float(d.get("regen_time_pct", 0.0)),
            regen_calls=int(d.get("regen_calls", 0)),
            rss_delta_mb=float(d.get("rss_delta_mb", 0.0)),
            rss_mb=float(d.get("rss_mb", 0.0) or 0.0),
            se_mean=float(d.get("se_mean", 0.0) or 0.0),
            se_std=float(d.get("se_std", 0.0) or 0.0),
        )

    @classmethod
    def from_telemetry_snapshot(cls, snapshot: Dict[str, Any]) -> "PolicyState":
        """Create from telemetry snapshot dict."""
        # Handle regen_time_pct computation
        elapsed_ms = snapshot.get("elapsed_sec", 1.0) * 1000
        regen_time_ms = snapshot.get("regen_time_ms", 0.0)
        regen_time_pct = (regen_time_ms / elapsed_ms * 100) if elapsed_ms > 0 else 0.0

        return cls(
            tok_per_sec=float(snapshot.get("tok_per_sec", 0.0)),
            mean_step_ms=float(snapshot.get("mean_step_ms", 0.0)),
            p95_step_ms=float(snapshot.get("p95_step_ms", 0.0)),
            steps_in_window=int(snapshot.get("steps_in_window", 0)),
            cache_hit_rate=float(snapshot.get("cache_hit_rate", 0.0)),
            cache_misses=int(snapshot.get("cache_misses", 0)),
            cache_hits=int(snapshot.get("cache_hits", 0)),
            regen_hit_rate=float(snapshot.get("regen_hit_rate", 1.0)),
            regen_time_pct=regen_time_pct,
            regen_calls=int(snapshot.get("regen_calls", 0)),
            rss_delta_mb=0.0,  # Delta requires previous snapshot
            rss_mb=float(snapshot.get("rss_mb", 0.0) or 0.0),
            se_mean=float(snapshot.get("se_mean", 0.0) or 0.0),
            se_std=0.0,  # Not in individual snapshots
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict."""
        return asdict(self)


# =============================================================================
# PolicyAction: 4-knob action vector (WO-ECHO-POLICY-GOLDEN-01)
# =============================================================================

@dataclass
class PolicyAction:
    """
    4-knob action vector for runtime control.

    WO-ECHO-POLICY-GOLDEN-01: Reduced from 7 to 4 knobs.
    Removed knobs that cannot be wired at runtime:
    - sample_rate: hardcoded 0.1 in code
    - cache_size: no cache_size parameter exists
    - block_rows: burned into CDNA format at encoding time

    Each field is an index into the corresponding value array:
    - verify_policy: 0=always, 1=sampled, 2=trust_cached
    - cpp_kernel: 0=off, 1=on
    - fused_matmul: 0=off, 1=on
    - fibpi_mode: 0=round_robin, 1=se_bucket

    Action space: 3 × 2 × 2 × 2 = 24 combinations
    """
    verify_policy: int = 2      # 0=always, 1=sampled, 2=trust_cached
    cpp_kernel: int = 0         # 0=off, 1=on
    fused_matmul: int = 0       # 0=off, 1=on
    fibpi_mode: int = 0         # 0=round_robin, 1=se_bucket

    @property
    def verify_policy_str(self) -> str:
        """Get verify_policy as string for passing to functions."""
        return VERIFY_POLICY_VALUES[self.verify_policy]

    @property
    def fibpi_mode_str(self) -> str:
        """Get fibpi_mode as string."""
        return FIBPI_MODE_VALUES[self.fibpi_mode]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict with actual values (not indices)."""
        return {
            "verify_policy": VERIFY_POLICY_VALUES[self.verify_policy],
            "cpp_kernel": bool(self.cpp_kernel),
            "fused_matmul": bool(self.fused_matmul),
            "fibpi_mode": FIBPI_MODE_VALUES[self.fibpi_mode],
        }

    def to_indices_dict(self) -> Dict[str, int]:
        """Convert to dict with indices."""
        return asdict(self)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "PolicyAction":
        """Create from dict (with actual values or indices)."""
        # Handle both formats: actual values or indices
        if isinstance(d.get("verify_policy"), str):
            # Actual values
            verify_idx = VERIFY_POLICY_VALUES.index(d["verify_policy"])
            cpp_kernel = 1 if d.get("cpp_kernel") else 0
            fused_matmul = 1 if d.get("fused_matmul") else 0
            fibpi_mode = 1 if d.get("fibpi_mode") == "se_bucket" else 0
            return cls(
                verify_policy=verify_idx,
                cpp_kernel=cpp_kernel,
                fused_matmul=fused_matmul,
                fibpi_mode=fibpi_mode,
            )
        else:
            # Indices - filter out deprecated keys
            filtered = {k: v for k, v in d.items()
                       if k in ("verify_policy", "cpp_kernel", "fused_matmul", "fibpi_mode")}
            return cls(**filtered)

    @classmethod
    def from_index(cls, index: int) -> "PolicyAction":
        """
        Create action from flat index (0-23).

        Useful for exhaustive sweeps over the action space.
        """
        if not 0 <= index < TOTAL_ACTION_COMBINATIONS:
            raise ValueError(f"Index {index} out of range [0, {TOTAL_ACTION_COMBINATIONS})")

        # Decode: index = verify*8 + cpp*4 + fused*2 + fibpi
        verify_policy = index // 8
        remainder = index % 8
        cpp_kernel = remainder // 4
        remainder = remainder % 4
        fused_matmul = remainder // 2
        fibpi_mode = remainder % 2

        return cls(
            verify_policy=verify_policy,
            cpp_kernel=cpp_kernel,
            fused_matmul=fused_matmul,
            fibpi_mode=fibpi_mode,
        )

    def to_index(self) -> int:
        """Convert action to flat index (0-23)."""
        return (self.verify_policy * 8 +
                self.cpp_kernel * 4 +
                self.fused_matmul * 2 +
                self.fibpi_mode)

    @classmethod
    def default(cls) -> "PolicyAction":
        """Return default (conservative) action."""
        return cls(
            verify_policy=2,    # trust_cached
            cpp_kernel=0,       # off
            fused_matmul=0,     # off
            fibpi_mode=0,       # round_robin
        )

    @classmethod
    def baseline(cls) -> "PolicyAction":
        """
        Return baseline action for golden verification.

        This is the reference action against which policy actions are compared.
        """
        return cls(
            verify_policy=2,    # trust_cached (for speed)
            cpp_kernel=0,       # off
            fused_matmul=0,     # off
            fibpi_mode=0,       # round_robin
        )

    @classmethod
    def random(cls, rng: Optional[np.random.Generator] = None) -> "PolicyAction":
        """Generate random action for exploration."""
        if rng is None:
            rng = np.random.default_rng()
        return cls(
            verify_policy=rng.integers(0, 3),
            cpp_kernel=rng.integers(0, 2),
            fused_matmul=rng.integers(0, 2),
            fibpi_mode=rng.integers(0, 2),
        )

    @classmethod
    def all_actions(cls) -> List["PolicyAction"]:
        """Generate all 24 possible actions."""
        return [cls.from_index(i) for i in range(TOTAL_ACTION_COMBINATIONS)]


# =============================================================================
# PolicyMLP: Tiny policy network with multi-head output
# =============================================================================

class PolicyMLP:
    """
    Tiny policy network with multi-head output.

    WO-ECHO-POLICY-GOLDEN-01: Reduced to 4 heads (was 7).

    Architecture:
        Input: [14] normalized state features
        -> Linear(14, 32) + ReLU
        -> Linear(32, 32) + ReLU
        -> Multi-head output:
            verify_policy: Softmax(3)
            cpp_kernel: Softmax(2)
            fused_matmul: Softmax(2)
            fibpi_mode: Softmax(2)

    ~2000 parameters. Trainable on CPU.
    """

    def __init__(self, state_dim: int = 14, hidden_dim: int = 32):
        torch, nn = _import_torch()

        self.state_dim = state_dim
        self.hidden_dim = hidden_dim

        # Shared backbone
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )

        # Multi-head output (4 heads for 4 knobs)
        self.heads = nn.ModuleDict({
            "verify_policy": nn.Linear(hidden_dim, HEAD_DIMS["verify_policy"]),
            "cpp_kernel": nn.Linear(hidden_dim, HEAD_DIMS["cpp_kernel"]),
            "fused_matmul": nn.Linear(hidden_dim, HEAD_DIMS["fused_matmul"]),
            "fibpi_mode": nn.Linear(hidden_dim, HEAD_DIMS["fibpi_mode"]),
        })

    def parameters(self):
        """Return all parameters."""
        torch, nn = _import_torch()
        params = list(self.shared.parameters())
        for head in self.heads.values():
            params.extend(head.parameters())
        return params

    def to(self, device):
        """Move to device."""
        self.shared = self.shared.to(device)
        for name in self.heads:
            self.heads[name] = self.heads[name].to(device)
        return self

    def eval(self):
        """Set to eval mode."""
        self.shared.eval()
        for head in self.heads.values():
            head.eval()
        return self

    def train(self, mode: bool = True):
        """Set to train mode."""
        self.shared.train(mode)
        for head in self.heads.values():
            head.train(mode)
        return self

    def forward(self, state) -> Dict[str, Any]:
        """
        Forward pass.

        Args:
            state: torch.Tensor of shape [batch, 14] or [14]

        Returns:
            Dict mapping head name to softmax probabilities
        """
        torch, _ = _import_torch()

        # Handle single state
        if state.dim() == 1:
            state = state.unsqueeze(0)

        # Shared features
        h = self.shared(state)

        # Multi-head output with softmax
        return {
            name: torch.softmax(head(h), dim=-1)
            for name, head in self.heads.items()
        }

    def __call__(self, state):
        return self.forward(state)

    def state_dict(self) -> Dict[str, Any]:
        """Return state dict."""
        torch, _ = _import_torch()
        state = {
            "shared": self.shared.state_dict(),
            "heads": {name: head.state_dict() for name, head in self.heads.items()},
            "config": {
                "state_dim": self.state_dim,
                "hidden_dim": self.hidden_dim,
            },
        }
        return state

    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        """Load state dict."""
        self.shared.load_state_dict(state_dict["shared"])
        for name, head_state in state_dict["heads"].items():
            self.heads[name].load_state_dict(head_state)

    @classmethod
    def from_state_dict(cls, state_dict: Dict[str, Any]) -> "PolicyMLP":
        """Create from state dict."""
        config = state_dict["config"]
        model = cls(
            state_dim=config["state_dim"],
            hidden_dim=config["hidden_dim"],
        )
        model.load_state_dict(state_dict)
        return model

    def count_parameters(self) -> int:
        """Count total parameters."""
        return sum(p.numel() for p in self.parameters())


# =============================================================================
# Reward Function
# =============================================================================

def compute_reward(
    prev_state: PolicyState,
    action: PolicyAction,
    next_state: PolicyState,
    correctness_ok: bool,
) -> float:
    """
    Compute reward for a state transition.

    Args:
        prev_state: State before action
        action: Action taken
        next_state: State after action
        correctness_ok: Whether verification passed

    Returns:
        Reward value (higher is better)

    Reward components:
    - Correctness gate: -10.0 if verification fails
    - Throughput improvement: +/- based on tok/s delta
    - Latency penalty: -0.3 if p95 increased by >50%
    - Memory penalty: -0.2 if RSS increased by >100MB
    - Cache bonus: +0.1 if cache hit rate >90%
    """
    # Correctness gate: HUGE penalty if verification fails
    if not correctness_ok:
        return -10.0

    # Throughput improvement (clipped to [-1, 1])
    tps_delta = next_state.tok_per_sec - prev_state.tok_per_sec
    tps_reward = float(np.clip(tps_delta / 100.0, -1.0, 1.0))

    # Latency penalty
    latency_penalty = 0.0
    if prev_state.p95_step_ms > 0 and next_state.p95_step_ms > prev_state.p95_step_ms * 1.5:
        latency_penalty = 0.3

    # Memory penalty
    memory_penalty = 0.0
    if next_state.rss_delta_mb > 100.0:
        memory_penalty = 0.2

    # Cache bonus
    cache_bonus = 0.0
    if next_state.cache_hit_rate > 0.9:
        cache_bonus = 0.1

    return tps_reward - latency_penalty - memory_penalty + cache_bonus


# =============================================================================
# State Normalization
# =============================================================================

@dataclass
class StateNormalizer:
    """Normalizes states using running mean and std."""
    mean: np.ndarray = field(default_factory=lambda: np.zeros(14, dtype=np.float32))
    std: np.ndarray = field(default_factory=lambda: np.ones(14, dtype=np.float32))
    count: int = 0
    eps: float = 1e-8

    def update(self, state: PolicyState) -> None:
        """Update running statistics with new state."""
        x = state.to_array()
        self.count += 1

        # Welford's online algorithm
        delta = x - self.mean
        self.mean = self.mean + delta / self.count
        delta2 = x - self.mean
        # Running variance approximation
        if self.count > 1:
            self.std = np.sqrt(
                ((self.count - 2) * self.std**2 + delta * delta2) / (self.count - 1)
            )
        self.std = np.maximum(self.std, self.eps)

    def normalize(self, state: PolicyState) -> np.ndarray:
        """Normalize state to zero mean, unit variance."""
        x = state.to_array()
        return (x - self.mean) / (self.std + self.eps)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to dict."""
        return {
            "mean": self.mean.tolist(),
            "std": self.std.tolist(),
            "count": self.count,
            "eps": self.eps,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "StateNormalizer":
        """Deserialize from dict."""
        return cls(
            mean=np.array(d["mean"], dtype=np.float32),
            std=np.array(d["std"], dtype=np.float32),
            count=d["count"],
            eps=d.get("eps", 1e-8),
        )


# =============================================================================
# PolicyHarness: Main Interface
# =============================================================================

class PolicyHarness:
    """
    Main interface for policy-driven runtime control.

    Provides:
    - State extraction from telemetry
    - Action selection (greedy or epsilon-greedy)
    - Action application via context manager
    - Checkpoint save/load
    """

    def __init__(
        self,
        model: PolicyMLP,
        normalizer: StateNormalizer,
    ):
        self.model = model
        self.normalizer = normalizer
        self._prev_rss_mb: Optional[float] = None

    @classmethod
    def default(cls) -> "PolicyHarness":
        """Return harness with default (untrained) policy."""
        model = PolicyMLP(state_dim=14, hidden_dim=32)
        normalizer = StateNormalizer()
        return cls(model=model, normalizer=normalizer)

    @classmethod
    def from_checkpoint(cls, path: str) -> "PolicyHarness":
        """Load trained policy from checkpoint."""
        torch, _ = _import_torch()

        checkpoint = torch.load(path, map_location="cpu", weights_only=False)

        model = PolicyMLP.from_state_dict(checkpoint["model"])
        model.eval()

        normalizer = StateNormalizer.from_dict(checkpoint["normalizer"])

        return cls(model=model, normalizer=normalizer)

    def save_checkpoint(self, path: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Save policy to checkpoint."""
        torch, _ = _import_torch()

        checkpoint = {
            "model": self.model.state_dict(),
            "normalizer": self.normalizer.to_dict(),
            "schema": "echo.policy.v0",
            "work_order": "WO-ECHO-POLICY-HARNESS-01",
        }

        if extra:
            checkpoint.update(extra)

        Path(path).parent.mkdir(parents=True, exist_ok=True)
        torch.save(checkpoint, path)

    def extract_state(self, telem: "Telemetry") -> PolicyState:
        """
        Extract PolicyState from Telemetry object.

        Uses the global counters and current step times.
        """
        gc = telem._global_counters
        c = telem._counters

        # Compute metrics
        tokens = c.tokens
        elapsed = 1.0  # Default
        if telem._last_snapshot_time:
            import time
            elapsed = max(time.perf_counter() - telem._last_snapshot_time, 0.001)

        tok_per_sec = tokens / elapsed if elapsed > 0 else 0.0

        step_times = c.step_times_ms
        import statistics
        mean_step = statistics.mean(step_times) if step_times else 0.0
        p95_step = 0.0
        if len(step_times) >= 20:
            p95_step = statistics.quantiles(step_times, n=20)[18]
        elif step_times:
            p95_step = max(step_times)

        cache_total = c.cache_hits + c.cache_misses
        cache_hit_rate = c.cache_hits / cache_total if cache_total > 0 else 0.0

        regen_hit_rate = 1.0 - (c.regen_calls / c.steps) if c.steps > 0 else 1.0

        # Memory
        from .telemetry import get_rss_mb
        rss_mb = get_rss_mb() or 0.0
        rss_delta_mb = 0.0
        if self._prev_rss_mb is not None:
            rss_delta_mb = rss_mb - self._prev_rss_mb
        self._prev_rss_mb = rss_mb

        # Regen time percentage
        elapsed_ms = elapsed * 1000
        regen_time_pct = (c.regen_time_ms / elapsed_ms * 100) if elapsed_ms > 0 else 0.0

        # Se
        se_mean = statistics.mean(c.se_values) if c.se_values else 0.0
        se_std = statistics.stdev(c.se_values) if len(c.se_values) > 1 else 0.0

        return PolicyState(
            tok_per_sec=tok_per_sec,
            mean_step_ms=mean_step,
            p95_step_ms=p95_step,
            steps_in_window=c.steps,
            cache_hit_rate=cache_hit_rate,
            cache_misses=c.cache_misses,
            cache_hits=c.cache_hits,
            regen_hit_rate=regen_hit_rate,
            regen_time_pct=regen_time_pct,
            regen_calls=c.regen_calls,
            rss_delta_mb=rss_delta_mb,
            rss_mb=rss_mb,
            se_mean=se_mean,
            se_std=se_std,
        )

    def select_action(
        self,
        state: PolicyState,
        explore: bool = False,
        epsilon: float = 0.1,
        rng: Optional[np.random.Generator] = None,
    ) -> PolicyAction:
        """
        Select action for given state.

        Args:
            state: Current state
            explore: If True, use epsilon-greedy exploration
            epsilon: Exploration probability
            rng: Random generator for exploration

        Returns:
            Selected action
        """
        torch, _ = _import_torch()

        if rng is None:
            rng = np.random.default_rng()

        # Epsilon-greedy exploration
        if explore and rng.random() < epsilon:
            return PolicyAction.random(rng)

        # Normalize state
        state_norm = self.normalizer.normalize(state)
        state_tensor = torch.tensor(state_norm, dtype=torch.float32)

        # Forward pass
        with torch.no_grad():
            probs = self.model(state_tensor)

        # Sample from each head (or take argmax for greedy)
        # WO-ECHO-POLICY-GOLDEN-01: Only 4 heads now
        action = PolicyAction(
            verify_policy=int(probs["verify_policy"].argmax().item()),
            cpp_kernel=int(probs["cpp_kernel"].argmax().item()),
            fused_matmul=int(probs["fused_matmul"].argmax().item()),
            fibpi_mode=int(probs["fibpi_mode"].argmax().item()),
        )

        return action

    def select_action_with_logprob(
        self,
        state: PolicyState,
    ) -> Tuple[PolicyAction, float]:
        """
        Select action and return log probability (for training).

        Returns:
            (action, log_prob)
        """
        torch, _ = _import_torch()

        # Normalize state
        state_norm = self.normalizer.normalize(state)
        state_tensor = torch.tensor(state_norm, dtype=torch.float32)

        # Forward pass
        probs = self.model(state_tensor)

        # Sample from each head
        log_prob = 0.0
        action_indices = {}

        for name, p in probs.items():
            dist = torch.distributions.Categorical(p.squeeze(0))
            idx = dist.sample()
            log_prob += dist.log_prob(idx).item()
            action_indices[name] = int(idx.item())

        action = PolicyAction(**action_indices)
        return action, log_prob

    @contextmanager
    def apply_action(self, action: PolicyAction) -> Generator[Dict[str, Any], None, None]:
        """
        Context manager to apply knobs during decode.

        WO-ECHO-POLICY-GOLDEN-01: Fixed env var names and removed unwired knobs.

        Sets environment variables for the runtime to pick up,
        then restores original values on exit.

        Note: verify_policy is NOT set via env var - it should be passed
        as a function parameter to stream_transformer_block_forward().
        Use action.verify_policy_str to get the string value.

        Env vars set:
        - HELIX_USE_CPP_KERNEL: "1" or "0"
        - HELIX_USE_FUSED_MATMUL: "1" or "0"
        - HELIX_FIBPI_MODE: "round_robin" or "se_bucket"

        Yields the action dict for logging.
        """
        action_dict = action.to_dict()

        # Map action to environment variables
        # WO-ECHO-POLICY-GOLDEN-01: Fixed env var names
        # - HELIX_CPP_KERNEL → HELIX_USE_CPP_KERNEL
        # - HELIX_FUSED_MATMUL → HELIX_USE_FUSED_MATMUL
        # - verify_policy NOT set via env (pass as parameter)
        # - Removed: sample_rate, cache_size, block_rows (unwirable)
        env_mapping = {
            "HELIX_USE_CPP_KERNEL": "1" if action_dict["cpp_kernel"] else "0",
            "HELIX_USE_FUSED_MATMUL": "1" if action_dict["fused_matmul"] else "0",
            "HELIX_FIBPI_MODE": action_dict["fibpi_mode"],
        }

        # Save original values
        original_env = {}
        for key in env_mapping:
            original_env[key] = os.environ.get(key)

        try:
            # Set new values
            for key, value in env_mapping.items():
                os.environ[key] = str(value)

            # WO-ECHO-POLICY-GOLDEN-BITE-01: Reload kernel config after env change
            # This updates CPP_OK/FUSED_MATMUL_OK module globals so dispatch uses new values
            try:
                from helix_cdc.regrow.stream_xw_matmul import reload_kernel_config
                reload_kernel_config()
            except ImportError:
                pass  # Module not available, skip reload

            yield action_dict

        finally:
            # Restore original values
            for key, original in original_env.items():
                if original is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = original

            # WO-ECHO-POLICY-GOLDEN-BITE-01: Reload again to restore original kernel config
            try:
                from helix_cdc.regrow.stream_xw_matmul import reload_kernel_config
                reload_kernel_config()
            except ImportError:
                pass


# =============================================================================
# Self-test
# =============================================================================

if __name__ == "__main__":
    import sys

    print("=" * 60)
    print("WO-ECHO-POLICY-GOLDEN-01 - Policy Harness Self-Test")
    print("=" * 60)
    print()

    # Test PolicyState
    print("1. PolicyState")
    state = PolicyState(
        tok_per_sec=42.5,
        mean_step_ms=23.4,
        p95_step_ms=45.6,
        steps_in_window=100,
        cache_hit_rate=0.85,
        cache_misses=15,
        cache_hits=85,
        regen_hit_rate=0.95,
        regen_time_pct=5.0,
        regen_calls=5,
        rss_delta_mb=10.0,
        rss_mb=500.0,
        se_mean=0.3,
        se_std=0.05,
    )
    arr = state.to_array()
    print(f"   to_array shape: {arr.shape}")
    print(f"   First 3 values: {arr[:3]}")

    # Test PolicyAction
    print("\n2. PolicyAction (reduced to 4 knobs)")
    action = PolicyAction.default()
    print(f"   Default action: {action.to_dict()}")
    print(f"   Action index: {action.to_index()}")
    action_random = PolicyAction.random()
    print(f"   Random action: {action_random.to_dict()}")

    # Test action space enumeration
    print(f"\n   Action space size: {TOTAL_ACTION_COMBINATIONS}")
    all_actions = PolicyAction.all_actions()
    print(f"   All actions generated: {len(all_actions)}")

    # Test from_index roundtrip
    for i in [0, 5, 10, 23]:
        a = PolicyAction.from_index(i)
        assert a.to_index() == i, f"Roundtrip failed for index {i}"
    print(f"   from_index/to_index roundtrip: PASS")

    # Test PolicyMLP
    print("\n3. PolicyMLP")
    model = PolicyMLP(state_dim=14, hidden_dim=32)
    print(f"   Parameters: {model.count_parameters()}")

    torch, _ = _import_torch()
    state_tensor = state.to_tensor()
    probs = model(state_tensor)
    print(f"   Output heads: {list(probs.keys())}")
    print(f"   verify_policy probs: {probs['verify_policy'].detach().numpy().round(3)}")

    # Test reward function
    print("\n4. compute_reward")
    next_state = PolicyState(
        tok_per_sec=50.0,  # Improved!
        mean_step_ms=20.0,
        p95_step_ms=40.0,
        steps_in_window=100,
        cache_hit_rate=0.92,  # Good cache
        cache_misses=8,
        cache_hits=92,
        regen_hit_rate=0.96,
        regen_time_pct=4.0,
        regen_calls=4,
        rss_delta_mb=5.0,
        rss_mb=505.0,
        se_mean=0.32,
        se_std=0.04,
    )
    reward = compute_reward(state, action, next_state, correctness_ok=True)
    print(f"   Reward (correctness OK): {reward:.4f}")
    reward_fail = compute_reward(state, action, next_state, correctness_ok=False)
    print(f"   Reward (correctness FAIL): {reward_fail:.4f}")

    # Test PolicyHarness
    print("\n5. PolicyHarness")
    harness = PolicyHarness.default()
    print(f"   Created default harness")

    # Test apply_action context manager
    print("\n6. apply_action context manager")
    action_test = PolicyAction(verify_policy=0, cpp_kernel=1, fused_matmul=1, fibpi_mode=1)
    with harness.apply_action(action_test) as applied:
        print(f"   Applied: {applied}")
        # WO-ECHO-POLICY-GOLDEN-01: Fixed env var names
        print(f"   HELIX_USE_CPP_KERNEL = {os.environ.get('HELIX_USE_CPP_KERNEL')}")
        print(f"   HELIX_USE_FUSED_MATMUL = {os.environ.get('HELIX_USE_FUSED_MATMUL')}")
        print(f"   HELIX_FIBPI_MODE = {os.environ.get('HELIX_FIBPI_MODE')}")
        assert os.environ.get('HELIX_USE_CPP_KERNEL') == '1', "CPP kernel not set!"
        assert os.environ.get('HELIX_USE_FUSED_MATMUL') == '1', "Fused matmul not set!"
        assert os.environ.get('HELIX_FIBPI_MODE') == 'se_bucket', "FIBPI mode not set!"

    # Verify env restored
    assert os.environ.get("HELIX_USE_CPP_KERNEL") is None, "Env not restored!"
    print("   Env restored after context exit: PASS")

    # Test save/load
    print("\n7. Checkpoint save/load")
    test_path = "/tmp/policy_harness_selftest.pt"
    harness.save_checkpoint(test_path)
    print(f"   Saved to {test_path}")

    harness2 = PolicyHarness.from_checkpoint(test_path)
    print(f"   Loaded from checkpoint")

    print()
    print("=" * 60)
    print("Self-test PASSED")
    print("=" * 60)
