"""
src/bandit_allocator.py

Contextual Bandit & Reinforcement Learning Dynamic Resource Allocator for Q-TensorFormer.
Replaces reactive PID loops with an online-learning LinUCB / Contextual Policy that maps
the 8D token information state z_t directly to the optimal Pareto decision tuple
(TT-rank, KV-precision, early-exit).
"""

from typing import Dict, Any, List, Tuple, Optional
import math
import torch
import numpy as np


class ContextualBanditAllocator:
    """
    LinUCB (Linear Upper Confidence Bound) Contextual Bandit Allocator.
    
    Observes 8D context vector z_t in [0, 1]^8 and selects from discrete action space:
    - Action 0: r=1,  INT4 KV, Early Exit Allowed  (Ultra-Edge)
    - Action 1: r=2,  INT4 KV, No Exit            (Edge)
    - Action 2: r=4,  INT8 KV, No Exit            (Moderate)
    - Action 3: r=8,  INT8 KV, No Exit            (Balanced)
    - Action 4: r=16, FP16 KV, No Exit            (High Capacity)
    """

    def __init__(self, context_dim: int = 8, n_actions: int = 5, alpha: float = 0.5):
        self.context_dim = context_dim
        self.n_actions = n_actions
        self.alpha = alpha  # Exploration parameter

        # LinUCB state parameters: A_a in R^{d x d}, b_a in R^{d}
        self.A = [np.identity(context_dim) for _ in range(n_actions)]
        self.b = [np.zeros((context_dim, 1)) for _ in range(n_actions)]

        # Action mapping
        self.action_specs = [
            {"rank": 1, "kv_precision": "INT4", "allow_exit": True,  "cost_multiplier": 0.05},
            {"rank": 2, "kv_precision": "INT4", "allow_exit": False, "cost_multiplier": 0.15},
            {"rank": 4, "kv_precision": "INT8", "allow_exit": False, "cost_multiplier": 0.35},
            {"rank": 8, "kv_precision": "INT8", "allow_exit": False, "cost_multiplier": 0.65},
            {"rank": 16, "kv_precision": "FP16", "allow_exit": False, "cost_multiplier": 1.00},
        ]

        # Statistics
        self.step_count = 0
        self.cumulative_reward = 0.0
        self.action_counts = [0] * n_actions

    def select_action(self, z_t: torch.Tensor) -> Tuple[int, Dict[str, Any]]:
        """
        Selects an action given 8D context vector z_t via LinUCB.
        
        Args:
            z_t: Tensor of shape (context_dim,) or (1, context_dim)
        """
        z = z_t.detach().cpu().numpy().flatten().reshape(-1, 1)
        p = np.zeros(self.n_actions)

        for a in range(self.n_actions):
            A_inv = np.linalg.inv(self.A[a])
            theta_a = A_inv @ self.b[a]
            # UCB score = expected reward + confidence bound
            expected_r = float(np.asarray(theta_a.T @ z).item())
            conf_bound = float(self.alpha * np.sqrt(max(0.0, float(np.asarray(z.T @ A_inv @ z).item()))))
            p[a] = expected_r + conf_bound

        chosen_action = int(np.argmax(p))
        self.action_counts[chosen_action] += 1
        self.step_count += 1

        spec = self.action_specs[chosen_action]
        telemetry = {
            "step": self.step_count,
            "chosen_action": chosen_action,
            "chosen_rank": spec["rank"],
            "chosen_kv_precision": spec["kv_precision"],
            "allow_exit": spec["allow_exit"],
            "predicted_score": float(p[chosen_action]),
            "action_counts": list(self.action_counts),
        }
        return chosen_action, telemetry

    def update(self, z_t: torch.Tensor, action: int, reward: float):
        """
        Updates LinUCB parameters after observing the actual reward.
        R_t = Delta_Q - lambda_L * Latency - lambda_M * Memory
        """
        z = z_t.detach().cpu().numpy().flatten().reshape(-1, 1)
        self.A[action] += z @ z.T
        self.b[action] += reward * z
        self.cumulative_reward += reward

    @staticmethod
    def compute_reward(
        token_surprise: float,
        action_rank: int,
        latency_budget_ratio: float = 1.0,
        memory_budget_ratio: float = 1.0,
    ) -> float:
        """
        Multi-objective reward function balancing prediction quality vs hardware pressure.
        """
        # Quality gain saturates if surprise is high and rank is high
        quality_gain = min(1.0, (action_rank / 16.0) * (token_surprise / 2.0))
        # Resource cost penalty
        resource_penalty = (action_rank / 16.0) * (0.5 * latency_budget_ratio + 0.5 * memory_budget_ratio)
        return float(quality_gain - resource_penalty)