"""
helix_cdc/routing/qubo_route_builder.py
=======================================

Build QUBO for unified route selection.

Binary Variables:
    x_i ∈ {0,1}  for i ∈ [0, K_max)     # Specialist selection
    y_j ∈ {0,1}  for j ∈ {cpu,gpu,qpu}  # Backend selection (exactly 1)
    a_k ∈ {0,1}  for k ∈ {B, C}         # World selection (exactly 1)

Objective (Minimize):
    E = λ_lat * Σⱼ(latency_j * yⱼ)           # Latency cost
      + λ_mem * Σᵢ(memory_i * xᵢ)            # Memory cost
      + λ_fid * fidelity_penalty(a, x)        # Fidelity deviation
      + λ_se  * (Se_actual - Se_target)²      # Se target deviation
      + λ_det * determinism_risk(a)           # Determinism penalty

      + A * (Σᵢxᵢ - K)²                       # Constraint: exactly K specialists
      + B * (Σⱼyⱼ - 1)²                       # Constraint: exactly 1 backend
      + C * (Σₖaₖ - 1)²                       # Constraint: exactly 1 world
"""

from __future__ import annotations

import hashlib
import json
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, field

import numpy as np

from .cost_extractors import (
    SpecialistCost,
    BackendCost,
    TelemetryCosts,
    extract_specialist_costs,
    extract_backend_costs,
    extract_telemetry_costs,
    compute_specialist_relevance,
    compute_fidelity_penalty,
    compute_determinism_risk,
    BACKEND_PROFILES,
    FIDELITY_BANDS,
)


@dataclass
class RouteQUBOConfig:
    """Configuration for QUBO route builder."""

    # Constraint penalties (Lagrangian multipliers)
    penalty_k_specialist: float = 100.0  # Exactly K specialists
    penalty_one_backend: float = 100.0   # Exactly 1 backend
    penalty_one_world: float = 100.0     # Exactly 1 world

    # Cost weights (lambdas)
    lambda_latency: float = 1.0
    lambda_memory: float = 1.0
    lambda_fidelity: float = 5.0
    lambda_se: float = 2.0
    lambda_determinism: float = 1.0
    lambda_thermal: float = 2.0

    # Se guardrails (hard thresholds)
    se_qpu_min: float = 0.15
    se_gpu_min: float = 0.10
    guardrail_penalty: float = 50.0

    # Hardware unavailable penalty (effectively infinite)
    hardware_offline_penalty: float = 1000.0

    # Memory budget
    vram_budget_mb: int = 6000
    memory_budget_mb: int = 8000


@dataclass
class QUBOMetadata:
    """Metadata about the constructed QUBO."""
    n_vars: int
    n_specialists: int
    n_backends: int = 3  # cpu, gpu, qpu
    n_worlds: int = 2    # world_b, world_c

    specialist_ids: List[str] = field(default_factory=list)
    backend_ids: List[str] = field(default_factory=lambda: ["cpu", "gpu", "qpu"])
    world_ids: List[str] = field(default_factory=lambda: ["world_b", "world_c"])

    # Variable index ranges
    specialist_range: Tuple[int, int] = (0, 0)
    backend_range: Tuple[int, int] = (0, 0)
    world_range: Tuple[int, int] = (0, 0)

    # Build parameters
    k_target: int = 5
    se_target: Optional[float] = None
    fidelity_band: str = "hybrid_cdna"
    vram_budget_mb: int = 6000

    # Hash for reproducibility
    qubo_sha256: str = ""


def build_route_qubo(
    request: Dict[str, Any],
    k_specialists: int = 5,
    fidelity_band: str = "hybrid_cdna",
    vram_budget_mb: int = 6000,
    se_target: Optional[float] = None,
    specialist_ids: Optional[List[str]] = None,
    specialist_costs: Optional[List[SpecialistCost]] = None,
    hardware_status: Optional[Dict[str, bool]] = None,
    telemetry: Optional[TelemetryCosts] = None,
    config: Optional[RouteQUBOConfig] = None,
) -> Tuple[Dict[Tuple[int, int], float], QUBOMetadata]:
    """
    Build QUBO for route selection.

    Args:
        request: Job request dict with problem type, shape, etc.
        k_specialists: Target number of specialists to select
        fidelity_band: "hybrid_cdna" | "world_b" | "lossless"
        vram_budget_mb: VRAM budget in MB
        se_target: Optional target Se value
        specialist_ids: Optional list of specialist IDs
        specialist_costs: Optional pre-computed specialist costs
        hardware_status: Dict with cpu_online, gpu_online, qpu_online
        telemetry: Current hardware telemetry
        config: RouteQUBOConfig for weights and penalties

    Returns:
        Q: QUBO dict {(i,j): weight} for upper triangle
        metadata: QUBOMetadata with variable mapping
    """
    if config is None:
        config = RouteQUBOConfig()

    # Get current telemetry if not provided
    if telemetry is None:
        telemetry = extract_telemetry_costs()

    # Default hardware status
    if hardware_status is None:
        hardware_status = {
            "cpu_online": True,
            "gpu_online": telemetry.gpu_available,
            "qpu_online": True,  # Simulated always available
        }

    # Generate specialist IDs if not provided
    if specialist_ids is None:
        # Default: generate from problem type
        problem_type = request.get("problem", {}).get("type", "inference")
        n_specialists = min(k_specialists * 2, 20)  # 2x candidates
        specialist_ids = [f"specialist_{i}" for i in range(n_specialists)]

    # Compute specialist costs if not provided
    if specialist_costs is None:
        specialist_costs = extract_specialist_costs(
            specialist_ids=specialist_ids,
            se_estimates=request.get("se_estimates"),
            chunk_stats=request.get("chunk_stats"),
        )

    # Get backend costs
    backend_costs = extract_backend_costs(
        hardware_status=hardware_status,
        telemetry=telemetry,
    )

    # Variable layout:
    # [0, n_specialists): x_i specialist selection
    # [n_specialists, n_specialists+3): y_j backend selection (cpu, gpu, qpu)
    # [n_specialists+3, n_specialists+5): a_k world selection (B, C)

    n_specialists = len(specialist_ids)
    n_backends = 3
    n_worlds = 2
    n_vars = n_specialists + n_backends + n_worlds

    specialist_range = (0, n_specialists)
    backend_range = (n_specialists, n_specialists + n_backends)
    world_range = (n_specialists + n_backends, n_vars)

    # Initialize QUBO as dict (upper triangle)
    Q: Dict[Tuple[int, int], float] = {}

    def add_linear(i: int, cost: float):
        """Add linear term Q[i,i] += cost"""
        key = (i, i)
        Q[key] = Q.get(key, 0.0) + cost

    def add_quadratic(i: int, j: int, cost: float):
        """Add quadratic term Q[i,j] += cost (upper triangle)"""
        if i > j:
            i, j = j, i
        key = (i, j)
        Q[key] = Q.get(key, 0.0) + cost

    # ========== CONSTRAINTS ==========

    # Constraint: Exactly K specialists
    # A * (Σᵢxᵢ - K)² = A * (Σᵢxᵢ² - 2K*Σᵢxᵢ + K² + 2*Σᵢ<ⱼ xᵢxⱼ)
    # Since xᵢ² = xᵢ: = A * ((1-2K)*Σᵢxᵢ + K² + 2*Σᵢ<ⱼ xᵢxⱼ)
    # Linear: (1-2K)*A on each xᵢ
    # Quadratic: 2A on each xᵢxⱼ pair
    A = config.penalty_k_specialist
    for i in range(specialist_range[0], specialist_range[1]):
        add_linear(i, A * (1 - 2 * k_specialists))
    for i in range(specialist_range[0], specialist_range[1]):
        for j in range(i + 1, specialist_range[1]):
            add_quadratic(i, j, 2 * A)

    # Constraint: Exactly 1 backend
    # B * (Σⱼyⱼ - 1)² = B * ((1-2)*Σⱼyⱼ + 1 + 2*Σⱼ<ₖ yⱼyₖ)
    #                 = B * (-Σⱼyⱼ + 2*Σⱼ<ₖ yⱼyₖ + 1)
    # Linear: -B on each yⱼ
    # Quadratic: 2B on each yⱼyₖ pair
    B = config.penalty_one_backend
    for j in range(backend_range[0], backend_range[1]):
        add_linear(j, -B)
    for j in range(backend_range[0], backend_range[1]):
        for k in range(j + 1, backend_range[1]):
            add_quadratic(j, k, 2 * B)

    # Constraint: Exactly 1 world
    C = config.penalty_one_world
    for k in range(world_range[0], world_range[1]):
        add_linear(k, -C)
    for k in range(world_range[0], world_range[1]):
        for l in range(k + 1, world_range[1]):
            add_quadratic(k, l, 2 * C)

    # ========== COST TERMS ==========

    backend_ids = ["cpu", "gpu", "qpu"]
    world_ids = ["world_b", "world_c"]

    # Latency costs: λ_lat * Σⱼ(latency_j * yⱼ)
    for j_offset, backend_id in enumerate(backend_ids):
        j = backend_range[0] + j_offset
        cost = backend_costs[backend_id]
        latency_normalized = cost.latency_base_us / 1000.0  # Normalize to ms scale
        add_linear(j, config.lambda_latency * latency_normalized)

    # Memory costs: λ_mem * Σᵢ(memory_i * xᵢ)
    for i_offset, spec_cost in enumerate(specialist_costs):
        i = specialist_range[0] + i_offset
        memory_normalized = spec_cost.memory_mb / 1000.0  # Normalize to GB scale
        add_linear(i, config.lambda_memory * memory_normalized)

    # Thermal costs on backends
    for j_offset, backend_id in enumerate(backend_ids):
        j = backend_range[0] + j_offset
        cost = backend_costs[backend_id]
        add_linear(j, config.lambda_thermal * cost.thermal_risk)

    # Hardware offline penalties
    for j_offset, backend_id in enumerate(backend_ids):
        j = backend_range[0] + j_offset
        cost = backend_costs[backend_id]
        if cost.availability < 0.5:  # Effectively offline
            add_linear(j, config.hardware_offline_penalty)

    # Se-based guardrails on backends
    # Estimate Se from request or specialists
    se_estimate = se_target
    if se_estimate is None:
        # Average Se from specialists
        se_values = [sc.se_score for sc in specialist_costs]
        se_estimate = sum(se_values) / len(se_values) if se_values else 0.5

    # GPU guardrail: penalty if Se < min
    if se_estimate < config.se_gpu_min:
        j_gpu = backend_range[0] + 1
        add_linear(j_gpu, config.guardrail_penalty)

    # QPU guardrail: penalty if Se < min
    if se_estimate < config.se_qpu_min:
        j_qpu = backend_range[0] + 2
        add_linear(j_qpu, config.guardrail_penalty)

    # Fidelity costs on specialists (depends on fidelity band)
    for i_offset, spec_cost in enumerate(specialist_costs):
        i = specialist_range[0] + i_offset
        fid_penalty = compute_fidelity_penalty(fidelity_band, spec_cost)
        add_linear(i, config.lambda_fidelity * fid_penalty)

    # Determinism risk on worlds
    for k_offset, world_id in enumerate(world_ids):
        k = world_range[0] + k_offset
        det_risk = compute_determinism_risk(world_id)
        add_linear(k, config.lambda_determinism * det_risk)

    # Se target deviation: λ_se * (Se_actual - Se_target)²
    # Approximation: penalize specialists far from target
    if se_target is not None:
        for i_offset, spec_cost in enumerate(specialist_costs):
            i = specialist_range[0] + i_offset
            relevance = compute_specialist_relevance(spec_cost, se_target)
            # Lower relevance = higher penalty (inverse)
            se_penalty = (1.0 - relevance) * config.lambda_se
            add_linear(i, se_penalty)

    # VRAM budget constraint (soft)
    # If GPU selected, specialists should fit in VRAM
    j_gpu = backend_range[0] + 1
    total_memory = sum(sc.memory_mb for sc in specialist_costs)
    if total_memory > vram_budget_mb:
        # Add coupling: selecting GPU with high-memory specialists is penalized
        memory_overflow_factor = (total_memory - vram_budget_mb) / vram_budget_mb
        for i_offset, spec_cost in enumerate(specialist_costs):
            i = specialist_range[0] + i_offset
            if spec_cost.memory_mb > vram_budget_mb / k_specialists:
                # This specialist exceeds average budget - penalize with GPU
                add_quadratic(i, j_gpu, memory_overflow_factor * 10.0)

    # Prefer world_c (CDNA) when fidelity_band is not world_b
    if fidelity_band != "world_b":
        k_world_b = world_range[0]
        add_linear(k_world_b, 20.0)  # Penalize world_b when fidelity required

    # ========== METADATA ==========

    metadata = QUBOMetadata(
        n_vars=n_vars,
        n_specialists=n_specialists,
        n_backends=n_backends,
        n_worlds=n_worlds,
        specialist_ids=specialist_ids,
        backend_ids=backend_ids,
        world_ids=world_ids,
        specialist_range=specialist_range,
        backend_range=backend_range,
        world_range=world_range,
        k_target=k_specialists,
        se_target=se_target,
        fidelity_band=fidelity_band,
        vram_budget_mb=vram_budget_mb,
    )

    # Compute QUBO hash for reproducibility
    qubo_canonical = json.dumps(
        {f"{i},{j}": v for (i, j), v in sorted(Q.items())},
        sort_keys=True,
        separators=(",", ":"),
    )
    metadata.qubo_sha256 = hashlib.sha256(qubo_canonical.encode()).hexdigest()

    return Q, metadata


def qubo_to_matrix(Q: Dict[Tuple[int, int], float], n: int) -> np.ndarray:
    """Convert QUBO dict to symmetric matrix."""
    matrix = np.zeros((n, n))
    for (i, j), v in Q.items():
        matrix[i, j] = v
        if i != j:
            matrix[j, i] = v
    return matrix


def qubo_to_bqm_format(
    Q: Dict[Tuple[int, int], float]
) -> Tuple[Dict[int, float], Dict[Tuple[int, int], float]]:
    """
    Convert QUBO dict to linear/quadratic format for BQM solvers.

    Returns:
        linear: {i: h_i}
        quadratic: {(i,j): J_ij} for i < j
    """
    linear = {}
    quadratic = {}

    for (i, j), v in Q.items():
        if i == j:
            linear[i] = linear.get(i, 0.0) + v
        else:
            key = (min(i, j), max(i, j))
            quadratic[key] = quadratic.get(key, 0.0) + v

    return linear, quadratic
