#!/usr/bin/env python3
"""
QUBO Router - Replaces threshold-based routing with optimization.

Instead of:
    if se >= 0.85: backend = 'qpu'
    elif se >= 0.45: backend = 'gpu'
    else: backend = 'cpu'

We now formulate a QUBO (Quadratic Unconstrained Binary Optimization):
    minimize E = A(sum(x)-1)^2 + latency_cost + thermal_cost + uncertainty_cost

Binary variables:
    x_c, x_g, x_q in {0,1} for CPU, GPU, QPU

This is real physics-based decision making using Se as ONE input among many.
"""

import numpy as np
from typing import Dict, Tuple, Optional, NamedTuple
from dataclasses import dataclass
import time

from .telemetry import TelemetrySnapshot, GPUMetrics, CPUMetrics, MemoryMetrics


@dataclass
class QUBOConfig:
    """Configuration for QUBO cost weights."""
    # Constraint penalty (must pick exactly one backend)
    constraint_penalty: float = 100.0

    # Latency weights (lower = prefer this backend for speed)
    latency_cpu: float = 1.0    # CPU is baseline
    latency_gpu: float = 0.5    # GPU is faster for parallel work
    latency_qpu: float = 0.3    # QPU is fastest for QUBO-native problems

    # Thermal risk weights (higher = more thermal concern)
    thermal_cpu_weight: float = 0.5
    thermal_gpu_weight: float = 1.0   # GPU heats up most
    thermal_qpu_weight: float = 0.3   # QPU runs cool (simulated)

    # Weight multipliers for objective terms
    lambda_latency: float = 1.0
    lambda_thermal: float = 2.0    # Prioritize not overheating

    # --- Se-SHAPED PENALTIES (ChatGPT fix) ---
    # These create preference curves so Se actually matters
    # QPU penalty = se_qpu_k * (1 - se)^se_qpu_p (large when Se LOW)
    # GPU penalty = se_gpu_k * (1 - se)^se_gpu_p (medium curve)
    # CPU penalty = se_cpu_k * se^se_cpu_p (low curve, prefers low Se)
    se_qpu_k: float = 10.0       # QPU penalty coefficient
    se_qpu_p: float = 2.0        # QPU penalty exponent (quadratic)
    se_gpu_k: float = 3.0        # GPU penalty coefficient
    se_gpu_p: float = 1.5        # GPU penalty exponent
    se_cpu_k: float = 5.0        # CPU penalty coefficient (for high Se)
    se_cpu_p: float = 2.0        # CPU penalty exponent

    # --- HARD GUARDRAILS ---
    # Forbid QPU below this Se threshold (unless CPU/GPU unavailable)
    se_qpu_min: float = 0.7      # QPU only for Se >= 0.7
    se_gpu_min: float = 0.3      # GPU only for Se >= 0.3
    guardrail_penalty: float = 50.0  # Soft guardrail penalty (added to cost)


@dataclass
class QUBOResult:
    """Result of QUBO optimization."""
    backend: str                    # 'cpu', 'gpu', or 'qpu'
    backend_idx: int               # 0=cpu, 1=gpu, 2=qpu
    energy: float                  # Objective value at solution
    all_energies: Dict[str, float] # Energy for each backend option
    solve_method: str              # 'brute_force' or 'qpu'
    solve_time_ms: float
    objective_breakdown: Dict[str, float]  # Individual cost terms


class QUBORouter:
    """
    QUBO-based backend router using Se + telemetry.

    This replaces hardcoded thresholds with optimization that considers:
    - Symbolic Entropy (Se): task complexity/uncertainty
    - Thermal state: GPU temp, power headroom
    - Memory pressure: available RAM/VRAM
    - Load: CPU load average

    The QUBO is small (3 binary vars) so we solve by brute force.
    Real QPU integration can be added later for larger scheduling problems.
    """

    def __init__(self, config: QUBOConfig = None):
        self.config = config or QUBOConfig()
        self.backends = ['cpu', 'gpu', 'qpu']
        self._qpu_available = False  # Set True if real QPU access exists

    def build_qubo(
        self,
        se_norm: float,
        telemetry: TelemetrySnapshot,
        queue_depth: int = 0,
        deadline_ms: Optional[float] = None,
        hardware_status: Optional[Dict[str, bool]] = None
    ) -> np.ndarray:
        """
        Build QUBO matrix Q where E = x^T Q x.

        For 3 binary variables [x_c, x_g, x_q], Q is 3x3.
        Diagonal Q[i,i] = linear cost for selecting backend i
        Off-diagonal Q[i,j] = interaction cost (used for constraint)

        The constraint "pick exactly one" becomes:
            A * (x_c + x_g + x_q - 1)^2
            = A * (x_c^2 + x_g^2 + x_q^2 + 2*x_c*x_g + 2*x_c*x_q + 2*x_g*x_q
                   - 2*x_c - 2*x_g - 2*x_q + 1)

        Since x_i is binary, x_i^2 = x_i, so:
            = A * (x_c + x_g + x_q + 2*(x_c*x_g + x_c*x_q + x_g*x_q) - 2*x_c - 2*x_g - 2*x_q + 1)
            = A * (-x_c - x_g - x_q + 2*(x_c*x_g + x_c*x_q + x_g*x_q) + 1)

        For QUBO form (ignoring constant):
            Diagonal: Q[i,i] += -A (from linear terms)
            Off-diag: Q[i,j] += 2A (from quadratic interaction)
        """
        cfg = self.config
        A = cfg.constraint_penalty

        # Initialize 3x3 QUBO matrix
        Q = np.zeros((3, 3))

        # --- Constraint: pick exactly one ---
        # This adds penalty if more or fewer than one backend is selected
        for i in range(3):
            Q[i, i] += -A  # Linear term
        for i in range(3):
            for j in range(i + 1, 3):
                Q[i, j] += 2 * A  # Off-diagonal interaction

        # --- Latency cost ---
        # Lower latency = lower cost = prefer this backend
        latency_costs = [cfg.latency_cpu, cfg.latency_gpu, cfg.latency_qpu]

        # Adjust for deadline pressure (if deadline is tight, prefer faster backends)
        if deadline_ms is not None and deadline_ms < 100:
            deadline_factor = 100 / max(deadline_ms, 1)  # Tighter deadline = higher factor
            latency_costs = [c * deadline_factor for c in latency_costs]

        for i, cost in enumerate(latency_costs):
            Q[i, i] += cfg.lambda_latency * cost

        # --- Thermal risk cost ---
        # Higher temp/power = higher cost for that backend
        thermal_costs = self._compute_thermal_costs(telemetry)
        for i, cost in enumerate(thermal_costs):
            Q[i, i] += cfg.lambda_thermal * cost

        # --- Se-SHAPED PENALTIES (ChatGPT fix) ---
        # These create preference curves so Se actually matters:
        # - QPU penalty LARGE when Se LOW (don't waste QPU on simple tasks)
        # - GPU penalty MEDIUM when Se LOW
        # - CPU penalty LARGE when Se HIGH (CPU struggles with complex tasks)
        se_penalties = self._compute_se_penalties(se_norm)
        for i, pen in enumerate(se_penalties):
            Q[i, i] += pen

        # --- HARD GUARDRAILS ---
        # Forbid backends below their Se thresholds (availability-aware)
        guardrails = self._compute_guardrails(se_norm, hardware_status)
        for i, pen in enumerate(guardrails):
            Q[i, i] += pen

        # --- Queue depth penalty ---
        # If queue is deep, prefer faster backends
        if queue_depth > 5:
            queue_factor = min(2.0, queue_depth / 5)
            Q[0, 0] += queue_factor * 0.5  # Penalty for slow CPU
            Q[1, 1] -= queue_factor * 0.2  # Slight bonus for GPU
            Q[2, 2] -= queue_factor * 0.3  # Bonus for QPU

        return Q

    def _compute_thermal_costs(self, telemetry: TelemetrySnapshot) -> Tuple[float, float, float]:
        """
        Compute thermal risk costs for each backend based on current temps.
        """
        cfg = self.config

        # CPU thermal cost
        cpu_temp = telemetry.cpu.temp_c
        cpu_load_norm = telemetry.cpu.load_1m / telemetry.cpu.core_count
        if cpu_temp is not None and cpu_temp > 70:
            cpu_thermal = cfg.thermal_cpu_weight * (cpu_temp - 70) / 30  # Scale 70-100C to 0-1
        else:
            cpu_thermal = cfg.thermal_cpu_weight * max(0, cpu_load_norm - 0.8)

        # GPU thermal cost
        if telemetry.gpu.available:
            gpu_temp = telemetry.gpu.temp_c
            gpu_power_ratio = telemetry.gpu.power_w / max(telemetry.gpu.power_limit_w, 1)

            # High temp or high power = high thermal risk
            gpu_thermal = cfg.thermal_gpu_weight * (
                max(0, (gpu_temp - 60) / 40) +  # Scale 60-100C to 0-1
                max(0, gpu_power_ratio - 0.7)   # Penalty above 70% power
            )
        else:
            # GPU not available = infinite thermal cost (won't be selected)
            gpu_thermal = 100.0

        # QPU thermal cost (simulated as low/constant for now)
        qpu_thermal = cfg.thermal_qpu_weight * 0.1

        return (cpu_thermal, gpu_thermal, qpu_thermal)

    def _compute_se_penalties(self, se_norm: float) -> Tuple[float, float, float]:
        """
        Compute Se-shaped penalties per ChatGPT's recommendation.

        These create preference curves so Se actually drives routing:
        - CPU penalty = k * se^p (LARGE when Se HIGH - CPU struggles with complexity)
        - GPU penalty = k * (1-se)^p (MEDIUM curve)
        - QPU penalty = k * (1-se)^p (LARGE when Se LOW - don't waste QPU)
        """
        cfg = self.config

        # CPU: penalize at HIGH Se (CPU struggles with complex tasks)
        cpu_penalty = cfg.se_cpu_k * (se_norm ** cfg.se_cpu_p)

        # GPU: moderate penalty at LOW Se (prefers medium complexity)
        gpu_penalty = cfg.se_gpu_k * ((1.0 - se_norm) ** cfg.se_gpu_p)

        # QPU: heavy penalty at LOW Se (reserve for quantum-hard)
        qpu_penalty = cfg.se_qpu_k * ((1.0 - se_norm) ** cfg.se_qpu_p)

        return (cpu_penalty, gpu_penalty, qpu_penalty)

    def _compute_guardrails(
        self,
        se_norm: float,
        hardware_status: Optional[Dict[str, bool]] = None
    ) -> Tuple[float, float, float]:
        """
        Compute hard guardrail penalties with last-viable-backend relaxation.

        Logic:
        1. Compute base guardrail penalties (forbid if Se below threshold)
        2. Check if all backends would be effectively forbidden
        3. Only relax the least-bad viable backend if needed

        This prevents "dead zones" while avoiding over-relaxation.
        """
        cfg = self.config

        # Check availability
        if hardware_status is None:
            hardware_status = {'cpu_online': True, 'gpu_online': True, 'qpu_online': True}

        cpu_online = hardware_status.get('cpu_online', True)
        gpu_online = hardware_status.get('gpu_online', True)
        qpu_online = hardware_status.get('qpu_online', True)

        # Count available backends
        available_count = sum([cpu_online, gpu_online, qpu_online])

        # Emergency mode: ONLY when exactly 1 backend available (per ChatGPT review)
        emergency_mode = (available_count == 1)

        # Degraded mode: 2 backends available (smaller relaxations only if needed)
        degraded_mode = (available_count == 2)

        # --- Step 1: Compute base guardrail penalties ---
        cpu_guardrail = 0.0  # CPU has no minimum Se requirement

        # GPU base guardrail
        if se_norm < cfg.se_gpu_min:
            gpu_guardrail = cfg.guardrail_penalty
        else:
            gpu_guardrail = 0.0

        # QPU base guardrail
        if se_norm < cfg.se_qpu_min:
            qpu_guardrail = cfg.guardrail_penalty
        else:
            qpu_guardrail = 0.0

        # --- Step 2: Check viability ---
        # A backend is "viable" if it's online AND (not heavily penalized OR is last resort)
        # Offline backends get +1000 penalty later, so they're never viable

        cpu_viable = cpu_online  # CPU always viable if online (no Se min)
        gpu_viable = gpu_online and (se_norm >= cfg.se_gpu_min)
        qpu_viable = qpu_online and (se_norm >= cfg.se_qpu_min)

        viable_count = sum([cpu_viable, gpu_viable, qpu_viable])

        # --- Step 3: Last-viable-backend relaxation ---
        # Only relax if NO backends would be viable without relaxation

        if viable_count == 0 and available_count > 0:
            # All viable backends are blocked by Se guardrails
            # Relax the "least bad" online backend

            if cpu_online:
                # CPU is always the safest fallback (no relaxation needed, it has no guardrail)
                pass
            elif gpu_online and not qpu_online:
                # Only GPU available, relax its guardrail
                gpu_guardrail = cfg.guardrail_penalty * 0.1  # Heavy relaxation
            elif qpu_online and not gpu_online:
                # Only QPU available, relax its guardrail
                qpu_guardrail = cfg.guardrail_penalty * 0.1  # Heavy relaxation
            elif gpu_online and qpu_online:
                # Both GPU and QPU online but CPU offline
                # Prefer GPU (it has lower Se threshold)
                gpu_guardrail = cfg.guardrail_penalty * 0.2

        # Emergency mode: single backend, must use it
        if emergency_mode:
            if gpu_online and not cpu_online and not qpu_online:
                gpu_guardrail = cfg.guardrail_penalty * 0.05  # Minimal penalty
            elif qpu_online and not cpu_online and not gpu_online:
                qpu_guardrail = cfg.guardrail_penalty * 0.05  # Minimal penalty

        return (cpu_guardrail, gpu_guardrail, qpu_guardrail)

    def solve_brute_force(self, Q: np.ndarray) -> Tuple[int, float, Dict[str, float]]:
        """
        Solve QUBO by brute force enumeration.

        For 3 binary variables, there are only 8 combinations.
        We evaluate all and return the minimum.
        """
        min_energy = float('inf')
        min_idx = 0
        all_energies = {}

        # Enumerate all 8 combinations
        for i in range(8):
            x = np.array([(i >> j) & 1 for j in range(3)], dtype=float)
            energy = float(x @ Q @ x)

            # Store named energy
            combo_name = "_".join(
                self.backends[j] for j in range(3) if x[j] == 1
            ) or "none"
            all_energies[combo_name] = energy

            if energy < min_energy:
                min_energy = energy
                min_idx = i

        # Decode winner
        winner_x = np.array([(min_idx >> j) & 1 for j in range(3)])
        winner_idx = np.argmax(winner_x) if winner_x.sum() == 1 else 0

        return winner_idx, min_energy, all_energies

    def solve(
        self,
        Q: np.ndarray,
        use_qpu: bool = False
    ) -> Tuple[int, float, Dict[str, float], str]:
        """
        Solve QUBO, optionally using QPU if available.

        Returns: (backend_idx, energy, all_energies, solve_method)
        """
        if use_qpu and self._qpu_available:
            # Future: call real QPU solver (D-Wave, etc.)
            # For now, fall back to brute force
            pass

        idx, energy, all_energies = self.solve_brute_force(Q)
        return idx, energy, all_energies, "brute_force"

    def route(
        self,
        se_norm: float,
        telemetry: TelemetrySnapshot,
        queue_depth: int = 0,
        deadline_ms: Optional[float] = None,
        hardware_status: Dict[str, bool] = None
    ) -> QUBOResult:
        """
        Main routing interface: compute optimal backend given current state.

        Args:
            se_norm: Normalized symbolic entropy (0-1)
            telemetry: Current telemetry snapshot
            queue_depth: Number of pending tasks
            deadline_ms: Optional deadline pressure
            hardware_status: Dict with 'cpu_online', 'gpu_online', 'qpu_online'

        Returns:
            QUBOResult with chosen backend and diagnostics
        """
        start_time = time.time()

        # Build QUBO (pass hardware_status for availability-aware guardrails)
        Q = self.build_qubo(se_norm, telemetry, queue_depth, deadline_ms, hardware_status)

        # Apply hardware availability as infinite penalty (hard forbid)
        # This dominates all other penalties (guardrails=50, Se=~10)
        if hardware_status:
            if not hardware_status.get('cpu_online', True):
                Q[0, 0] += 1000.0  # CPU unavailable
            if not hardware_status.get('gpu_online', True):
                Q[1, 1] += 1000.0  # GPU unavailable
            if not hardware_status.get('qpu_online', True):
                Q[2, 2] += 1000.0  # QPU unavailable

        # Solve
        backend_idx, energy, all_energies, method = self.solve(Q)

        solve_time_ms = (time.time() - start_time) * 1000

        # Compute objective breakdown for explainability
        thermal_costs = self._compute_thermal_costs(telemetry)
        se_penalties = self._compute_se_penalties(se_norm)
        guardrails = self._compute_guardrails(se_norm, hardware_status)
        breakdown = {
            "se_norm": se_norm,
            "thermal_cpu": thermal_costs[0],
            "thermal_gpu": thermal_costs[1],
            "thermal_qpu": thermal_costs[2],
            "se_penalty_cpu": se_penalties[0],
            "se_penalty_gpu": se_penalties[1],
            "se_penalty_qpu": se_penalties[2],
            "guardrail_cpu": guardrails[0],
            "guardrail_gpu": guardrails[1],
            "guardrail_qpu": guardrails[2],
            "queue_depth": queue_depth,
            "deadline_ms": deadline_ms,
        }

        return QUBOResult(
            backend=self.backends[backend_idx],
            backend_idx=backend_idx,
            energy=energy,
            all_energies=all_energies,
            solve_method=method,
            solve_time_ms=solve_time_ms,
            objective_breakdown=breakdown
        )


# Module-level singleton
qubo_router = QUBORouter()


def route_with_qubo(
    se_norm: float,
    telemetry: TelemetrySnapshot,
    queue_depth: int = 0,
    deadline_ms: Optional[float] = None,
    hardware_status: Dict[str, bool] = None
) -> QUBOResult:
    """Convenience function for QUBO-based routing."""
    return qubo_router.route(
        se_norm=se_norm,
        telemetry=telemetry,
        queue_depth=queue_depth,
        deadline_ms=deadline_ms,
        hardware_status=hardware_status
    )


if __name__ == "__main__":
    # Quick test
    import json
    from .telemetry import read_telemetry

    print("Testing QUBO Router...")

    # Get current telemetry
    telem = read_telemetry()
    print(f"Telemetry: GPU={telem.gpu.available}, CPU load={telem.cpu.load_1m:.2f}")

    # Test routing at different Se levels
    for se in [0.2, 0.5, 0.7, 0.9]:
        result = route_with_qubo(
            se_norm=se,
            telemetry=telem,
            queue_depth=3,
            deadline_ms=500
        )
        print(f"Se={se:.1f} -> {result.backend} (energy={result.energy:.2f}, {result.solve_time_ms:.2f}ms)")
        print(f"  All energies: {result.all_energies}")
