src/resource_allocator.py

"""
Information-Value Resource Allocator Module for Q-TensorFormer.

The Intellectual Core:
  Instead of asking "How difficult is this token?", the allocator asks:
  "Given token information state z_t, hardware device state H_device, and current
   budgets, what is the cheapest additional computation that produces the greatest
   expected marginal improvement?"

Action Space:
  - TT Rank: 1, 2, 4, 8
  - Attention Pathway: classical_fast (SDPA), classical_standard, quantum_qksam
  - Computation Depth: skip, partial, full
  - KV Cache Precision: FP16, INT8, INT4, evict

Includes:
  - MarginalValueModel: empirical neural predictor of Delta Q and Delta Costs
  - InformationValueAllocator: dimensionally consistent marginal utility routing
  - Hysteresis & Anti-Chattering stabilization
  - Routing Churn & Stability tracking
  - PIDDualSubgradientController: closed-loop empirical SLA tracking with diagnostics
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, Optional, Tuple, List, NamedTuple, Any, Union
from dataclasses import dataclass, field


class AllocatorAction(NamedTuple):
    rank: int              # 1, 2, 4, 8
    attention_mode: str    # "classical_fast", "classical_standard", "quantum_qksam"
    depth_mode: str        # "skip", "partial", "full"
    kv_precision: str      # "fp16", "int8", "int4", "evict"
    kv_residency: str = "hot_gpu"  # "hot_gpu", "warm_cpu", "cold_evicted"


@dataclass
class AllocationBudget:
    max_latency_ms: Optional[float] = None
    max_memory_mb: Optional[float] = None
    max_peak_memory_mb: Optional[float] = None
    max_energy_uj: Optional[float] = None
    max_energy_per_token_j: Optional[float] = None
    max_kv_mb: Optional[float] = None
    max_ttft_ms: Optional[float] = None
    max_tpot_ms: Optional[float] = None
    max_bandwidth_gb_s: Optional[float] = None
    max_cost_usd_1m: Optional[float] = None
    min_quality_target: float = 0.0
    min_quality_fidelity: float = 0.90
    risk_tolerance: float = 0.5
    phase: str = "decode"                 # "prefill" or "decode"
    workload_type: str = "general"        # "reasoning", "math", "code", "dialogue", "long_context"
    lambda_latency: float = 1.0
    lambda_memory: float = 0.5
    lambda_energy: float = 0.2
    lambda_bandwidth: float = 0.3
    lambda_cost: float = 0.1


class MarginalValueModel(nn.Module):
    """
    Learned Marginal-Value Predictor for Q-TensorFormer.

    Predicts expected quality gain (Delta Q) and hardware resource increments
    (Delta latency, Delta memory, Delta energy, Delta bandwidth) for candidate
    actions conditioned on token information state z_t, hardware state, and budget state.

    Mathematical Formulation:
      Delta Q_hat, Delta C_hat = f_theta(z_t, a, H_device, B_state)
    """

    CANDIDATE_RANKS = [1, 2, 4, 8]
    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]
    DEPTH_MODES = ["skip", "partial", "full"]
    KV_MODES = ["fp16", "int8", "int4"]

    def __init__(self, info_dim: int = 8, hidden_dim: int = 64):
        super().__init__()
        self.info_dim = info_dim
        self.hidden_dim = hidden_dim

        # Action encoding: rank (4), attention (3), depth (3), kv (3) = 13 dims
        self.action_dim = 4 + 3 + 3 + 3
        self.hw_dim = 3       # latency_pressure, memory_pressure, bandwidth_pressure
        self.budget_dim = 4   # lambda_l, lambda_m, lambda_e, lambda_b

        in_dim = info_dim + self.action_dim + self.hw_dim + self.budget_dim

        self.backbone = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
        )

        # 7 output heads: Quality, Latency, Memory, Energy, Bandwidth, Financial Cost, Epistemic Uncertainty
        self.head_quality = nn.Linear(hidden_dim, 1)      # Delta Q in [0, 1]
        self.head_latency = nn.Linear(hidden_dim, 1)      # Delta Latency in ms >= 0
        self.head_memory = nn.Linear(hidden_dim, 1)       # Delta Memory in MB >= 0
        self.head_energy = nn.Linear(hidden_dim, 1)       # Delta Energy in uJ >= 0
        self.head_bandwidth = nn.Linear(hidden_dim, 1)    # Delta Bandwidth in Bytes >= 0
        self.head_cost = nn.Linear(hidden_dim, 1)         # Delta Cost in $/1M tokens >= 0
        self.head_uncertainty = nn.Linear(hidden_dim, 1)  # Epistemic Uncertainty sigma in [0, 1]

    def encode_action(
        self,
        rank_idx: int,
        attn_idx: int,
        depth_idx: int,
        kv_idx: int,
        device: torch.device,
    ) -> torch.Tensor:
        """One-hot encodes the action components into a 13-dim vector."""
        vec = torch.zeros(self.action_dim, device=device)
        vec[rank_idx] = 1.0
        vec[4 + attn_idx] = 1.0
        vec[7 + depth_idx] = 1.0
        vec[10 + kv_idx] = 1.0
        return vec

    def forward(
        self,
        z_t: torch.Tensor,
        action_enc: torch.Tensor,
        hw_state: torch.Tensor,
        budget_state: torch.Tensor,
    ) -> Dict[str, torch.Tensor]:
        """
        Forward pass predicting marginal outcomes.

        Args:
            z_t: (B, T, 8) or (N, 8)
            action_enc: (B, T, 13) or (N, 13)
            hw_state: (B, T, 3) or (N, 3)
            budget_state: (B, T, 4) or (N, 4)

        Returns:
            Dict containing predicted delta_q, delta_latency_ms, delta_memory_mb,
            delta_energy_uj, delta_bandwidth_bytes, delta_cost_usd, uncertainty
        """
        x = torch.cat([z_t, action_enc, hw_state, budget_state], dim=-1)
        h = self.backbone(x)

        dq = torch.sigmoid(self.head_quality(h)).squeeze(-1)
        dlat = F.softplus(self.head_latency(h)).squeeze(-1) * 5.0
        dmem = F.softplus(self.head_memory(h)).squeeze(-1) * 2.0
        dnrg = F.softplus(self.head_energy(h)).squeeze(-1) * 5000.0
        dbw = F.softplus(self.head_bandwidth(h)).squeeze(-1) * 30000.0
        dcost = F.softplus(self.head_cost(h)).squeeze(-1) * 2.50
        dunc = torch.sigmoid(self.head_uncertainty(h)).squeeze(-1)

        return {
            "delta_q": dq,
            "delta_latency_ms": dlat,
            "delta_memory_mb": dmem,
            "delta_energy_uj": dnrg,
            "delta_bandwidth_bytes": dbw,
            "delta_cost_usd": dcost,
            "uncertainty": dunc,
        }

    def compute_marginal_utility(
        self,
        z_t: torch.Tensor,
        rank_idx: int,
        attn_idx: int,
        depth_idx: int,
        kv_idx: int,
        hw_state: torch.Tensor,
        budget_state: torch.Tensor,
        eps: float = 1e-4,
    ) -> torch.Tensor:
        """
        Compute dimensionally consistent marginal utility:
          Value(a | z_t) = Delta Q / (Delta C_eff + eps)
        """
        device = z_t.device
        act_enc = self.encode_action(rank_idx, attn_idx, depth_idx, kv_idx, device)
        # Expand act_enc, hw_state, budget_state to match z_t shape
        shape = z_t.shape[:-1]
        act_enc_expanded = act_enc.reshape(*([1] * len(shape)), self.action_dim).expand(*shape, self.action_dim)
        hw_expanded = hw_state.reshape(*([1] * len(shape)), self.hw_dim).expand(*shape, self.hw_dim)
        b_expanded = budget_state.reshape(*([1] * len(shape)), self.budget_dim).expand(*shape, self.budget_dim)

        preds = self.forward(z_t, act_enc_expanded, hw_expanded, b_expanded)

        # Dimensionless normalized cost combination
        # Normalize: latency (ms / 10ms), memory (MB / 2MB), energy (uJ / 20000uJ), bandwidth (Bytes / 65600B)
        norm_lat = preds["delta_latency_ms"] / 10.0
        norm_mem = preds["delta_memory_mb"] / 2.0
        norm_nrg = preds["delta_energy_uj"] / 20000.0
        norm_bw = preds["delta_bandwidth_bytes"] / 65600.0

        lambda_l = b_expanded[..., 0]
        lambda_m = b_expanded[..., 1]
        lambda_e = b_expanded[..., 2]
        lambda_b = b_expanded[..., 3]

        eff_cost = norm_lat * (1.0 + lambda_l) + norm_mem * lambda_m + norm_nrg * lambda_e + norm_bw * lambda_b
        utility = preds["delta_q"] / (eff_cost + eps)

        return utility


class InformationValueAllocator(nn.Module):
    """
    Closed-loop resource allocation controller with marginal utility modeling.
    """

    CANDIDATE_RANKS = [1, 2, 4, 8]
    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]
    DEPTH_MODES = ["skip", "partial", "full"]
    KV_MODES = ["fp16", "int8", "int4"]

    def __init__(
        self,
        info_dim: int = 8,
        hidden_dim: int = 32,
        hysteresis_tau: float = 0.15,
        default_preset: str = "balanced",
        marginal_value_model: Optional[MarginalValueModel] = None,
    ):
        super().__init__()
        self.info_dim = info_dim
        self.hysteresis_tau = hysteresis_tau
        self.default_preset = default_preset
        self.marginal_value_model = marginal_value_model

        # Learned quality gain estimator: predicts Delta Q(a | z_t) for each candidate rank
        self.quality_rank_net = nn.Sequential(
            nn.Linear(info_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(self.CANDIDATE_RANKS)),
        )

        # Learned gate for attention mode: [fast, standard, quantum]
        self.quality_attn_net = nn.Sequential(
            nn.Linear(info_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(self.ATTENTION_MODES)),
        )

        # Learned gate for depth: [skip, partial, full]
        self.quality_depth_net = nn.Sequential(
            nn.Linear(info_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(self.DEPTH_MODES)),
        )

        # Learned gate for KV precision: [fp16, int8, int4]
        self.quality_kv_net = nn.Sequential(
            nn.Linear(info_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(self.KV_MODES)),
        )

        # Inductive bias calibration:
        # High entropy (z[1]) and uncertainty (z[2]) scale quality gains for ranks 4 and 8
        with torch.no_grad():
            self.quality_rank_net[0].weight.data.normal_(0, 0.05)
            self.quality_rank_net[0].weight.data[0, 1] += 2.0  # H_t
            self.quality_rank_net[0].weight.data[1, 2] += 2.5  # U_t
            self.quality_rank_net[2].weight.data.normal_(0, 0.05)
            self.quality_rank_net[2].weight.data[2, 0] += 1.8  # rank 4
            self.quality_rank_net[2].weight.data[3, 0] += 2.8  # rank 8
            self.quality_rank_net[2].weight.data[3, 1] += 2.2  # rank 8 on uncertainty
            self.quality_rank_net[2].bias.data.copy_(torch.tensor([-0.2, 0.2, 0.6, 1.1]))

            # Quantum attention boosted when uncertainty is high
            self.quality_attn_net[0].weight.data[0, 2] += 3.0
            self.quality_attn_net[2].weight.data[2, 0] += 2.5

            # Base depth preference: [skip, partial, full]
            self.quality_depth_net[2].bias.data.copy_(torch.tensor([-2.0, 0.0, 2.0]))

        # Anti-chattering / Hysteresis state tracking
        self.register_buffer("prev_rank_idx", torch.tensor(2, dtype=torch.long))  # default rank 4 (idx 2)
        self.register_buffer("prev_attn_idx", torch.tensor(0, dtype=torch.long))  # default classical_fast
        self.register_buffer("total_decisions", torch.tensor(0, dtype=torch.long))
        self.register_buffer("churn_count", torch.tensor(0, dtype=torch.long))

        # Base nominal costs for actions (empirically normalized relative units)
        self.cost_ranks = [0.15, 0.30, 0.60, 1.00]      # ranks 1, 2, 4, 8
        self.cost_attn = [0.20, 0.50, 1.80]             # fast, standard, quantum
        self.cost_depth = [0.05, 0.40, 1.00]            # skip, partial, full
        self.cost_kv = [1.00, 0.50, 0.25]               # fp16, int8, int4

    def reset_stability_counters(self):
        """Reset churn and total decision counters."""
        self.total_decisions.zero_()
        self.churn_count.zero_()
        self.prev_rank_idx.fill_(2)
        self.prev_attn_idx.fill_(0)

    @property
    def routing_churn_rate(self) -> float:
        """Percentage of token steps where routing changed between successive steps."""
        tot = max(1, self.total_decisions.item())
        return self.churn_count.item() / tot

    def forward(
        self,
        z_t: torch.Tensor,
        budget: Optional[AllocationBudget] = None,
        preset: Optional[str] = None,
        force_classical: bool = False,
    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, float]]:
        """
        Evaluate marginal utility and select optimal action per token.
        """
        B, T, _ = z_t.shape
        device = z_t.device
        mode = (preset or self.default_preset).lower()

        # Extract weights from budget
        b = budget or AllocationBudget()
        lambda_l = b.lambda_latency
        lambda_m = b.lambda_memory
        lambda_e = b.lambda_energy
        lambda_b = b.lambda_bandwidth

        # Adjust lambda weights based on deployment preset
        if mode == "latency":
            lambda_l *= 2.5
        elif mode == "memory":
            lambda_m *= 3.0
        elif mode == "energy":
            lambda_e *= 3.0
        elif mode == "edge":
            lambda_l *= 2.0
            lambda_m *= 2.5
            lambda_e *= 2.5
            force_classical = True
        elif mode == "classical_only":
            force_classical = True
        elif mode == "full":
            lambda_l *= 0.2
            lambda_m *= 0.2
            lambda_e *= 0.2

        # 1. Rank Selection: Value(r | z_t) = Delta Q_r / (Cost_r * (1 + lambda_l * L + lambda_m * M) + eps)
        rank_logits = self.quality_rank_net(z_t)  # (B, T, 4)
        est_dq_rank = torch.sigmoid(rank_logits)

        # Resource pressure: z_t[..., 5]=L, z_t[..., 6]=M, z_t[..., 7]=B
        L_pressure = z_t[..., 5].unsqueeze(-1)
        M_pressure = z_t[..., 6].unsqueeze(-1)
        B_pressure = z_t[..., 7].unsqueeze(-1)
        E_pressure = (L_pressure + M_pressure) / 2.0

        rank_costs = torch.tensor(self.cost_ranks, device=device).reshape(1, 1, 4)
        cost_multiplier = 0.40 * (
            1.0 + lambda_l * L_pressure + lambda_m * M_pressure + lambda_b * B_pressure + lambda_e * E_pressure
        )
        effective_rank_cost = cost_multiplier * rank_costs
        utility_rank = est_dq_rank - effective_rank_cost  # (B, T, 4) Lagrangian dual objective

        # Per-token best rank indices
        token_best_rank_idx = torch.argmax(utility_rank, dim=-1)  # (B, T)

        # Hysteresis stabilization on sequence level
        mean_utility_rank = utility_rank.mean(dim=(0, 1))  # (4,)
        best_rank_idx = int(torch.argmax(mean_utility_rank).item())
        prev_idx = self.prev_rank_idx.item()

        delta_u = mean_utility_rank[best_rank_idx] - mean_utility_rank[prev_idx]
        if delta_u < self.hysteresis_tau:
            chosen_rank_idx = prev_idx
        else:
            chosen_rank_idx = best_rank_idx
            if self.training or not torch.is_grad_enabled():
                if chosen_rank_idx != prev_idx:
                    self.churn_count.add_(1)
            self.prev_rank_idx.fill_(chosen_rank_idx)

        self.total_decisions.add_(1)
        chosen_rank = self.CANDIDATE_RANKS[chosen_rank_idx]

        # 2. Attention Pathway Selection (Token-level granularity)
        attn_logits = self.quality_attn_net(z_t)  # (B, T, 3)
        est_dq_attn = torch.sigmoid(attn_logits)
        attn_costs = torch.tensor(self.cost_attn, device=device).reshape(1, 1, 3)
        cost_multiplier_attn = 0.35 * (1.0 + lambda_l * L_pressure + lambda_e * E_pressure)
        utility_attn = est_dq_attn - cost_multiplier_attn * attn_costs  # (B, T, 3)

        if force_classical:
            utility_attn[..., 2] = -float("inf")

        chosen_attn_idx = torch.argmax(utility_attn, dim=-1)  # (B, T)

        # 3. Depth Execution (Layer-level or sequence-level)
        depth_logits = self.quality_depth_net(z_t)  # (B, T, 3)
        depth_scores = torch.softmax(depth_logits, dim=-1)  # (B, T, 3)
        avg_depth = depth_scores.mean(dim=(0, 1))
        chosen_depth_idx = int(torch.argmax(avg_depth).item())
        chosen_depth = self.DEPTH_MODES[chosen_depth_idx]

        # 4. KV Cache Policy Selection
        kv_logits = self.quality_kv_net(z_t)  # (B, T, 3)
        kv_costs = torch.tensor(self.cost_kv, device=device).reshape(1, 1, 3)
        cost_multiplier_kv = 0.40 * (1.0 + lambda_m * M_pressure * 2.0)
        utility_kv = torch.sigmoid(kv_logits) - cost_multiplier_kv * kv_costs
        mean_kv_u = utility_kv.mean(dim=(0, 1))
        chosen_kv_idx = int(torch.argmax(mean_kv_u).item())
        chosen_kv = self.KV_MODES[chosen_kv_idx]

        # Calculate diagnostics
        q_routed_tokens = (chosen_attn_idx == 2).sum().item()
        total_tokens = B * T
        q_usage_pct = (q_routed_tokens / max(1, total_tokens)) * 100.0

        # Per-token ranks mapped
        token_ranks = torch.tensor(self.CANDIDATE_RANKS, device=device)[token_best_rank_idx]  # (B, T)

        diagnostics = {
            "chosen_rank": chosen_rank,
            "mean_rank": float(token_ranks.float().mean().item()),
            "quantum_usage_pct": round(q_usage_pct, 2),
            "chosen_depth": chosen_depth,
            "chosen_kv_precision": chosen_kv,
            "routing_churn_rate": round(self.routing_churn_rate, 4),
            "effective_rank_cost": round(effective_rank_cost.mean().item(), 3),
        }

        decisions = {
            "rank": chosen_rank,
            "token_ranks": token_ranks,                # (B, T) per-token rank
            "attn_mode_idx": chosen_attn_idx,          # (B, T)
            "depth_mode": chosen_depth,
            "kv_precision": chosen_kv,
            "is_quantum_token": (chosen_attn_idx == 2),  # (B, T) bool
        }

        return decisions, diagnostics


class PIDDualSubgradientController:
    """
    Online Closed-Loop Dual Multiplier Controller for Q-TensorFormer.

    Tunes Lagrange multipliers lambda_k for latency, memory, energy, and bandwidth
    to empirically track user-specified SLA targets.

    Empirically Verifiable Control Metrics:
      - Settling time (t_settle): tokens until error enters +/- 5% tolerance band
      - Maximum overshoot (M_p): peak violation percentage above target
      - Steady-state error (e_ss): mean absolute error in the terminal window
      - Violation rate: percentage of steps where measured > budget
    """

    def __init__(
        self,
        target_latency_ms: Optional[float] = None,
        target_memory_mb: Optional[float] = None,
        target_energy_uj: Optional[float] = None,
        target_bandwidth_bytes: Optional[float] = None,
        kp: float = 0.05,
        ki: float = 0.01,
        kd: float = 0.005,
        lambda_min: float = 0.05,
        lambda_max: float = 10.0,
    ):
        self.target_latency_ms = target_latency_ms
        self.target_memory_mb = target_memory_mb
        self.target_energy_uj = target_energy_uj
        self.target_bandwidth_bytes = target_bandwidth_bytes
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.lambda_min = lambda_min
        self.lambda_max = lambda_max

        # Current multiplier states
        self.lambda_latency = 1.0
        self.lambda_memory = 0.5
        self.lambda_energy = 0.2
        self.lambda_bandwidth = 0.3

        # Integrals and previous errors
        self.integral_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0}
        self.prev_errors = {"latency": 0.0, "memory": 0.0, "energy": 0.0, "bandwidth": 0.0}
        self.history: List[Dict[str, float]] = []

    def update(
        self,
        measured_latency_ms: Optional[float] = None,
        measured_memory_mb: Optional[float] = None,
        measured_energy_uj: Optional[float] = None,
        measured_bandwidth_bytes: Optional[float] = None,
    ) -> AllocationBudget:
        """
        Update dual multipliers given observed empirical hardware metrics.
        Returns an updated AllocationBudget.
        """
        err_lat = 0.0
        err_mem = 0.0
        err_nrg = 0.0
        err_bw = 0.0

        if self.target_latency_ms is not None and measured_latency_ms is not None:
            err_lat = measured_latency_ms - self.target_latency_ms
            self.integral_errors["latency"] = max(-5.0, min(5.0, self.integral_errors["latency"] + err_lat))
            deriv = err_lat - self.prev_errors["latency"]
            self.prev_errors["latency"] = err_lat
            delta = self.kp * err_lat + self.ki * self.integral_errors["latency"] + self.kd * deriv
            self.lambda_latency = max(self.lambda_min, min(self.lambda_max, self.lambda_latency + delta))

        if self.target_memory_mb is not None and measured_memory_mb is not None:
            err_mem = measured_memory_mb - self.target_memory_mb
            self.integral_errors["memory"] = max(-5.0, min(5.0, self.integral_errors["memory"] + err_mem))
            deriv = err_mem - self.prev_errors["memory"]
            self.prev_errors["memory"] = err_mem
            delta = self.kp * err_mem + self.ki * self.integral_errors["memory"] + self.kd * deriv
            self.lambda_memory = max(self.lambda_min, min(self.lambda_max, self.lambda_memory + delta))

        if self.target_energy_uj is not None and measured_energy_uj is not None:
            err_nrg = measured_energy_uj - self.target_energy_uj
            self.integral_errors["energy"] = max(-5.0, min(5.0, self.integral_errors["energy"] + err_nrg))
            deriv = err_nrg - self.prev_errors["energy"]
            self.prev_errors["energy"] = err_nrg
            delta = self.kp * err_nrg + self.ki * self.integral_errors["energy"] + self.kd * deriv
            self.lambda_energy = max(self.lambda_min, min(self.lambda_max, self.lambda_energy + delta))

        if self.target_bandwidth_bytes is not None and measured_bandwidth_bytes is not None:
            err_bw = (measured_bandwidth_bytes - self.target_bandwidth_bytes) / 1000.0
            self.integral_errors["bandwidth"] = max(-5.0, min(5.0, self.integral_errors["bandwidth"] + err_bw))
            deriv = err_bw - self.prev_errors["bandwidth"]
            self.prev_errors["bandwidth"] = err_bw
            delta = self.kp * err_bw + self.ki * self.integral_errors["bandwidth"] + self.kd * deriv
            self.lambda_bandwidth = max(self.lambda_min, min(self.lambda_max, self.lambda_bandwidth + delta))

        record = {
            "step": len(self.history) + 1,
            "measured_latency_ms": measured_latency_ms or 0.0,
            "measured_memory_mb": measured_memory_mb or 0.0,
            "measured_energy_uj": measured_energy_uj or 0.0,
            "measured_bandwidth_bytes": measured_bandwidth_bytes or 0.0,
            "error_latency": err_lat,
            "error_memory": err_mem,
            "error_energy": err_nrg,
            "error_bandwidth": err_bw,
            "lambda_latency": round(self.lambda_latency, 4),
            "lambda_memory": round(self.lambda_memory, 4),
            "lambda_energy": round(self.lambda_energy, 4),
            "lambda_bandwidth": round(self.lambda_bandwidth, 4),
        }
        self.history.append(record)

        return AllocationBudget(
            max_latency_ms=self.target_latency_ms,
            max_memory_mb=self.target_memory_mb,
            max_energy_uj=self.target_energy_uj,
            lambda_latency=self.lambda_latency,
            lambda_memory=self.lambda_memory,
            lambda_energy=self.lambda_energy,
            lambda_bandwidth=self.lambda_bandwidth,
        )

    def get_diagnostics(self) -> Dict[str, float]:
        """Compute control metrics over history."""
        if not self.history:
            return {}

        lat_errors = [h["error_latency"] for h in self.history if self.target_latency_ms is not None]
        if not lat_errors:
            return {"total_steps": len(self.history)}

        violations = sum(1 for e in lat_errors if e > 0)
        violation_rate = (violations / len(lat_errors)) * 100.0

        target = self.target_latency_ms or 1.0
        overshoot_pct = max(0.0, max(lat_errors) / target * 100.0)

        # Steady-state error over final 20%
        w = max(1, int(len(lat_errors) * 0.2))
        ss_error = sum(abs(e) for e in lat_errors[-w:]) / w

        # Settling time: step index where error remains within +/- 5% of target
        band = 0.05 * target
        settling_step = len(lat_errors)
        for i in range(len(lat_errors)):
            if all(abs(e) <= band for e in lat_errors[i:]):
                settling_step = i + 1
                break

        return {
            "total_steps": len(self.history),
            "violation_rate_pct": round(violation_rate, 2),
            "max_overshoot_pct": round(overshoot_pct, 2),
            "steady_state_error": round(ss_error, 4),
            "settling_step": settling_step,
        }

    def get_budget(self) -> AllocationBudget:
        return AllocationBudget(
            max_latency_ms=self.target_latency_ms,
            max_memory_mb=self.target_memory_mb,
            max_energy_uj=self.target_energy_uj,
            lambda_latency=self.lambda_latency,
            lambda_memory=self.lambda_memory,
            lambda_energy=self.lambda_energy,
            lambda_bandwidth=self.lambda_bandwidth,
        )


class ConstrainedDecisionEngine(nn.Module):
    """
    Unified Closed-Loop Constrained Decision Engine for Q-TensorFormer.

    Integrates:
      1. Information state evaluation (z_t)
      2. Multi-objective marginal value & resource cost prediction (Delta Q, Delta L, Delta M, Delta B, Delta E, Delta $)
      3. Risk-aware utility penalization with conservative fallback when confidence is low
      4. Binding constraint detection (identifies which SLA ceiling is throttling inference)
      5. Phase-aware execution modes (Prefill vs Decode)
      6. Workload-specific adaptation (Reasoning, Math, Code, Dialogue, Long-Context)
      7. Closed-loop PID feedback adaptation of Lagrangian shadow prices
    """

    CANDIDATE_RANKS = [1, 2, 4, 8]
    ATTENTION_MODES = ["classical_fast", "classical_standard", "quantum_qksam"]
    DEPTH_MODES = ["skip", "partial", "full"]
    KV_MODES = ["fp16", "int8", "int4"]
    KV_RESIDENCIES = ["hot_gpu", "warm_cpu", "cold_evicted"]

    def __init__(
        self,
        info_dim: int = 8,
        marginal_value_model: Optional[MarginalValueModel] = None,
        pid_controller: Optional[PIDDualSubgradientController] = None,
        risk_gamma: float = 0.35,
        uncertainty_threshold: float = 0.65,
        default_preset: str = "balanced",
    ):
        super().__init__()
        self.info_dim = info_dim
        self.marginal_value_model = marginal_value_model or MarginalValueModel(info_dim=info_dim)
        self.pid_controller = pid_controller or PIDDualSubgradientController()
        self.risk_gamma = risk_gamma
        self.uncertainty_threshold = uncertainty_threshold
        self.default_preset = default_preset

        # Historical state tracking
        self.last_action: Optional[AllocatorAction] = None
        self.step_count = 0
        self.binding_constraints_history: List[str] = []

    def evaluate_candidates(
        self,
        z_t: torch.Tensor,
        hw_state: Optional[torch.Tensor] = None,
        budget: Optional[AllocationBudget] = None,
    ) -> Tuple[AllocatorAction, Dict[str, Any]]:
        """
        Evaluates candidate actions under multi-budget constraints and risk penalties.
        Returns selected AllocatorAction and detailed diagnostics.
        """
        device = z_t.device
        budget = budget or self.pid_controller.get_budget()

        if hw_state is None:
            hw_state = torch.tensor([0.2, 0.3, 0.25], device=device)

        budget_vec = torch.tensor([
            budget.lambda_latency,
            budget.lambda_memory,
            budget.lambda_energy,
            budget.lambda_bandwidth,
        ], device=device)

        # Build candidate action space
        phase = getattr(budget, "phase", "decode").lower()
        candidates = []
        for r_idx, r in enumerate(self.CANDIDATE_RANKS):
            for a_idx, attn in enumerate(self.ATTENTION_MODES):
                for d_idx, depth in enumerate(self.DEPTH_MODES):
                    for k_idx, kv in enumerate(self.KV_MODES):
                        # In prefill phase, do not skip entire layer to preserve context representation
                        if phase == "prefill" and depth == "skip":
                            continue
                        candidates.append((r_idx, a_idx, d_idx, k_idx, AllocatorAction(r, attn, depth, kv, "hot_gpu")))

        best_action = None
        best_score = -float("inf")
        diagnostics: Dict[str, Any] = {}

        # Evaluate candidate utility
        for r_idx, a_idx, d_idx, k_idx, action in candidates:
            act_enc = self.marginal_value_model.encode_action(r_idx, a_idx, d_idx, k_idx, device)
            shape = z_t.shape[:-1]
            act_expanded = act_enc.reshape(*([1] * len(shape)), -1).expand(*shape, -1)
            hw_expanded = hw_state.reshape(*([1] * len(shape)), -1).expand(*shape, -1)
            b_expanded = budget_vec.reshape(*([1] * len(shape)), -1).expand(*shape, -1)

            preds = self.marginal_value_model(z_t, act_expanded, hw_expanded, b_expanded)

            dq = preds["delta_q"].mean().item()
            dlat = preds["delta_latency_ms"].mean().item()
            dmem = preds["delta_memory_mb"].mean().item()
            dnrg = preds["delta_energy_uj"].mean().item()
            dbw = preds["delta_bandwidth_bytes"].mean().item()
            dcost = preds["delta_cost_usd"].mean().item()
            dunc = preds["uncertainty"].mean().item()

            # Workload-specific inductive scaling
            workload = getattr(budget, "workload_type", "general").lower()
            if workload in ["reasoning", "math"]:
                if action.rank >= 4 and action.depth_mode == "full":
                    dq *= 1.25
            elif workload == "code":
                if action.rank >= 4:
                    dq *= 1.15
            elif workload == "dialogue":
                if action.rank <= 2 and action.kv_precision == "int4":
                    dlat *= 0.85

            # Dimensionless effective constraint penalty
            c_eff = (
                (dlat / 10.0) * budget.lambda_latency +
                (dmem / 2.0) * budget.lambda_memory +
                (dnrg / 20000.0) * budget.lambda_energy +
                (dbw / 65600.0) * budget.lambda_bandwidth +
                (dcost / 2.50) * getattr(budget, "lambda_cost", 0.1)
            )

            # Risk-penalized Lagrangian dual objective
            risk_penalty = self.risk_gamma * dunc
            score = dq - c_eff - risk_penalty

            # Check hard SLA limits if specified
            violates_budget = False
            if budget.max_tpot_ms is not None and dlat > budget.max_tpot_ms:
                violates_budget = True
            if budget.max_memory_mb is not None and dmem > budget.max_memory_mb:
                violates_budget = True
            if budget.max_energy_uj is not None and dnrg > budget.max_energy_uj:
                violates_budget = True

            if not violates_budget and score > best_score:
                best_score = score
                best_action = action
                diagnostics = {
                    "score": round(score, 4),
                    "expected_delta_q": round(dq, 4),
                    "expected_latency_ms": round(dlat, 3),
                    "expected_memory_mb": round(dmem, 3),
                    "expected_energy_uj": round(dnrg, 2),
                    "expected_bandwidth_bytes": int(dbw),
                    "expected_cost_usd": round(dcost, 4),
                    "uncertainty": round(dunc, 4),
                    "is_fallback": False,
                }

        # Safe fallback if high uncertainty or no feasible candidate found
        if best_action is None or diagnostics.get("uncertainty", 0.0) > self.uncertainty_threshold:
            best_action = AllocatorAction(rank=4, attention_mode="classical_standard", depth_mode="full", kv_precision="int8", kv_residency="hot_gpu")
            diagnostics["is_fallback"] = True
            diagnostics["fallback_reason"] = "uncertainty_exceeded" if best_action else "no_feasible_budget_candidate"

        # Detect binding constraint
        binding = "none"
        if budget.max_tpot_ms and diagnostics.get("expected_latency_ms", 0) >= 0.85 * budget.max_tpot_ms:
            binding = "latency"
        elif budget.max_memory_mb and diagnostics.get("expected_memory_mb", 0) >= 0.85 * budget.max_memory_mb:
            binding = "memory"
        elif budget.max_energy_uj and diagnostics.get("expected_energy_uj", 0) >= 0.85 * budget.max_energy_uj:
            binding = "energy"
        diagnostics["binding_constraint"] = binding
        self.binding_constraints_history.append(binding)

        self.last_action = best_action
        self.step_count += 1
        return best_action, diagnostics

    def step_feedback(
        self,
        measured_latency_ms: Optional[float] = None,
        measured_memory_mb: Optional[float] = None,
        measured_energy_uj: Optional[float] = None,
        measured_bandwidth_bytes: Optional[float] = None,
    ) -> AllocationBudget:
        """Closed-loop feedback update for PID dual subgradient multipliers."""
        return self.pid_controller.update(
            measured_latency_ms=measured_latency_ms,
            measured_memory_mb=measured_memory_mb,
            measured_energy_uj=measured_energy_uj,
            measured_bandwidth_bytes=measured_bandwidth_bytes,
        )

    def forward(
        self,
        z_t: torch.Tensor,
        budget: Optional[AllocationBudget] = None,
        hw_state: Optional[torch.Tensor] = None,
    ) -> Tuple[AllocatorAction, Dict[str, Any]]:
        """Forward pass delegating to evaluate_candidates for PyTorch module compliance."""
        return self.evaluate_candidates(z_t, hw_state=hw_state, budget=budget)