"""
helix_cdc/routing/cost_extractors.py
====================================

Extract costs from existing infrastructure for QUBO route builder.

Sources:
- ChunkStats: H/C/D metrics per specialist
- TelemetryObserver: Hardware state (GPU temp, memory, CPU load)
- GateController: Historical performance data
- CodecRouter: Lane-specific costs

All costs are normalized to [0, 1] for QUBO compatibility.
"""

from __future__ import annotations

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

import numpy as np


@dataclass
class SpecialistCost:
    """Cost profile for a single specialist (tensor/shard)."""
    specialist_id: str
    memory_mb: float = 0.0
    latency_ms: float = 0.0
    se_score: float = 0.5  # Symbolic entropy [0,1]
    fidelity_score: float = 1.0  # Quality [0,1], higher = better
    h_entropy: float = 0.5  # Shannon entropy [0,8] normalized
    c_coherence: float = 0.5  # Structure [0,1]
    d_depth: float = 0.5  # Complexity [0,1]


@dataclass
class BackendCost:
    """Cost profile for a compute backend."""
    backend_id: str
    latency_base_us: float = 50.0
    memory_overhead_mb: float = 0.0
    power_factor: float = 1.0  # Relative power consumption
    availability: float = 1.0  # 0 = offline, 1 = fully available
    thermal_risk: float = 0.0  # 0 = cool, 1 = overheating


@dataclass
class TelemetryCosts:
    """Current hardware state as costs."""
    gpu_temp_c: float = 50.0
    gpu_power_ratio: float = 0.5
    gpu_memory_used_mb: float = 0.0
    gpu_memory_total_mb: float = 8000.0
    gpu_available: bool = True
    cpu_load_norm: float = 0.3
    cpu_temp_c: Optional[float] = None
    memory_pressure_pct: float = 50.0
    timestamp_unix: float = field(default_factory=time.time)


# Backend latency profiles (from router_policy.py)
BACKEND_PROFILES = {
    "cpu": {
        "latency_us": 50.0,
        "memory_overhead_mb": 0.0,
        "power_factor": 0.3,
        "min_se": 0.0,  # CPU has no minimum Se
    },
    "gpu": {
        "latency_us": 20.0,
        "memory_overhead_mb": 500.0,
        "power_factor": 1.0,
        "min_se": 0.10,  # GPU preferred for Se >= 0.10
    },
    "qpu": {
        "latency_us": 800.0,  # Includes queueing/communication
        "memory_overhead_mb": 0.0,
        "power_factor": 0.1,  # QPU very power efficient
        "min_se": 0.15,  # QPU only for Se >= 0.15
    },
}

# Fidelity bands from CLAUDE.md
FIDELITY_BANDS = {
    "hybrid_cdna": {
        "max_error": 0.001,
        "world": "world_c",
        "deterministic": True,
        "receipt_flag": "FIDELITY: HYBRID_CDNA_VERIFIED",
    },
    "world_b": {
        "max_error": float("inf"),  # No fidelity guarantee
        "world": "world_b",
        "deterministic": True,  # Deterministic but not faithful
        "receipt_flag": "FIDELITY: WORLD_B_PRG",
    },
    "lossless": {
        "max_error": 0.0,
        "world": "world_c",
        "deterministic": True,
        "receipt_flag": "FIDELITY: LOSSLESS",
    },
}


def extract_specialist_costs(
    specialist_ids: List[str],
    se_estimates: Optional[Dict[str, float]] = None,
    chunk_stats: Optional[Dict[str, Dict]] = None,
) -> List[SpecialistCost]:
    """
    Extract cost profiles for specialists.

    Args:
        specialist_ids: List of specialist IDs (tensor names, shard IDs, etc.)
        se_estimates: Optional pre-computed Se values per specialist
        chunk_stats: Optional H/C/D stats from ChunkStats.compute()

    Returns:
        List of SpecialistCost objects
    """
    costs = []

    for sid in specialist_ids:
        cost = SpecialistCost(specialist_id=sid)

        # Apply Se estimate if provided
        if se_estimates and sid in se_estimates:
            cost.se_score = se_estimates[sid]

        # Apply chunk stats if provided (from codec_router)
        if chunk_stats and sid in chunk_stats:
            stats = chunk_stats[sid]
            cost.h_entropy = stats.get("H", 0.5)
            cost.c_coherence = stats.get("C", 0.5)
            cost.d_depth = stats.get("D", 0.5)

            # Derive Se from H*C*D (normalized)
            se_raw = cost.h_entropy * cost.c_coherence * cost.d_depth
            cost.se_score = min(1.0, se_raw / 4.0)  # Normalize assuming max raw ~4

        # Estimate memory from specialist type
        if "attn" in sid.lower():
            cost.memory_mb = 100.0  # Attention layers are larger
        elif "ffn" in sid.lower() or "mlp" in sid.lower():
            cost.memory_mb = 200.0  # FFN layers are largest
        else:
            cost.memory_mb = 50.0  # Default

        costs.append(cost)

    return costs


def extract_backend_costs(
    hardware_status: Optional[Dict[str, bool]] = None,
    telemetry: Optional[TelemetryCosts] = None,
) -> Dict[str, BackendCost]:
    """
    Extract cost profiles for backends based on hardware state.

    Args:
        hardware_status: Dict with 'cpu_online', 'gpu_online', 'qpu_online'
        telemetry: Current hardware telemetry

    Returns:
        Dict mapping backend_id to BackendCost
    """
    if hardware_status is None:
        hardware_status = {"cpu_online": True, "gpu_online": True, "qpu_online": True}

    costs = {}

    for backend_id, profile in BACKEND_PROFILES.items():
        online_key = f"{backend_id}_online"
        is_online = hardware_status.get(online_key, True)

        cost = BackendCost(
            backend_id=backend_id,
            latency_base_us=profile["latency_us"],
            power_factor=profile["power_factor"],
            availability=1.0 if is_online else 0.0,
        )

        # Apply telemetry-based adjustments
        if telemetry:
            if backend_id == "gpu":
                cost.availability = 1.0 if telemetry.gpu_available else 0.0

                # Thermal risk from GPU temp
                if telemetry.gpu_temp_c > 70:
                    cost.thermal_risk = min(1.0, (telemetry.gpu_temp_c - 70) / 30)
                else:
                    cost.thermal_risk = 0.0

                # Memory overhead
                cost.memory_overhead_mb = profile["memory_overhead_mb"]

            elif backend_id == "cpu":
                # CPU thermal risk from load
                cost.thermal_risk = max(0.0, telemetry.cpu_load_norm - 0.8)

            elif backend_id == "qpu":
                # QPU assumed cool (simulated)
                cost.thermal_risk = 0.1

        costs[backend_id] = cost

    return costs


def extract_telemetry_costs() -> TelemetryCosts:
    """
    Extract current hardware telemetry.

    Tries to use flowtorch.telemetry if available, otherwise returns defaults.

    Returns:
        TelemetryCosts with current hardware state
    """
    costs = TelemetryCosts()

    try:
        from flowtorch.telemetry import TelemetryObserver
        observer = TelemetryObserver()
        snapshot = observer.snapshot()

        # GPU metrics
        if snapshot.gpu.available:
            costs.gpu_temp_c = snapshot.gpu.temp_c
            costs.gpu_power_ratio = snapshot.gpu.power_w / max(snapshot.gpu.power_limit_w, 1)
            costs.gpu_memory_used_mb = snapshot.gpu.memory_used_mb
            costs.gpu_memory_total_mb = snapshot.gpu.memory_total_mb
            costs.gpu_available = True
        else:
            costs.gpu_available = False

        # CPU metrics
        costs.cpu_load_norm = snapshot.cpu.load_1m / max(snapshot.cpu.core_count, 1)
        costs.cpu_temp_c = snapshot.cpu.temp_c

        # Memory metrics
        costs.memory_pressure_pct = snapshot.memory.pressure_pct

        costs.timestamp_unix = snapshot.timestamp_unix

    except ImportError:
        # flowtorch not available, use defaults
        pass
    except Exception:
        # Telemetry collection failed, use defaults
        pass

    return costs


def compute_specialist_relevance(
    specialist_cost: SpecialistCost,
    se_target: Optional[float] = None,
) -> float:
    """
    Compute relevance score for a specialist given target Se.

    Specialists close to target Se are more relevant.

    Returns:
        Relevance score [0, 1], higher = more relevant
    """
    if se_target is None:
        # No target, all equally relevant
        return 0.5

    # Gaussian-like relevance around target
    distance = abs(specialist_cost.se_score - se_target)
    relevance = np.exp(-2.0 * distance ** 2)

    return float(relevance)


def compute_fidelity_penalty(
    fidelity_band: str,
    specialist_cost: SpecialistCost,
) -> float:
    """
    Compute fidelity penalty for specialist in given fidelity band.

    Higher fidelity bands penalize specialists with low quality scores.

    Returns:
        Penalty [0, inf), 0 = no penalty
    """
    band = FIDELITY_BANDS.get(fidelity_band, FIDELITY_BANDS["hybrid_cdna"])

    if band["max_error"] == float("inf"):
        # World B: no fidelity requirement
        return 0.0

    if band["max_error"] == 0.0:
        # Lossless: heavy penalty for any quality loss
        return (1.0 - specialist_cost.fidelity_score) * 100.0

    # Hybrid CDNA: moderate penalty
    return (1.0 - specialist_cost.fidelity_score) * 10.0


def compute_determinism_risk(world: str) -> float:
    """
    Compute determinism risk for world selection.

    World B (PRG) is deterministic but produces different outputs than original.

    Returns:
        Risk score [0, 1]
    """
    if world == "world_b":
        return 0.3  # Deterministic but not faithful
    elif world == "world_c":
        return 0.0  # Deterministic and faithful
    else:
        return 0.5  # Unknown
