#!/usr/bin/env python3
"""
FlowTorch Symbolic Entropy Router - Se-based Heterogeneous Compute Routing

Routes tensors to CPU/GPU/QPU based on Symbolic Entropy (Se) complexity metrics.
Integrates with the FlowBraid unified tensor bridge.

Se = H × U × D (REAL formula, not variance proxy):
    H = 1 - energy_at_10pct       # Entropy: spread-out spectrum = high H
    U = 1 - neighbor_coherence    # Unstructuredness: no spatial structure = high U
    D = sqrt(rank_ratio)          # Depth: high effective rank = high D

2D Routing (NEW 2026-01-26):
Uses (Se, C_struct) jointly where C_struct = max(row_coherence, col_coherence):
    Zone 1: Se < 0.30, C_struct >= 0.30 → CPU (simple + organized)
    Zone 2: 0.30 <= Se < 0.70, any      → GPU (parallelizable)
    Zone 3: Se >= 0.70, C_struct < 0.30 → QPU (complex + unstructured)
    Zone 4: Se >= 0.70, C_struct >= 0.30 → GPU (complex but structured, exploitable)

Now supports:
- Real tensor Se via tensor_se_estimator (default)
- 2D (Se, C_struct) routing for smarter backend selection
- QUBO-based routing (use_qubo_router=True) for physics-aware optimization
- Size-based gating (tensors > 80% VRAM → CPU)
"""

import numpy as np
import sys
from pathlib import Path
from typing import Optional, Dict

# Add project root for imports
PROJECT_ROOT = Path(__file__).parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

# Import real tensor Se estimator
try:
    from tools.tensor_se_estimator import (
        compute_tensor_se,
        compute_routing_decision,
        CPU_MAX_SE,
        GPU_MAX_SE,
        TENSOR_CLASSIFIER_AVAILABLE,
    )
    TENSOR_SE_AVAILABLE = True
except ImportError:
    TENSOR_SE_AVAILABLE = False
    CPU_MAX_SE = 0.30
    GPU_MAX_SE = 0.70
    TENSOR_CLASSIFIER_AVAILABLE = False

# Routing thresholds (from plan: CPU < 0.30, GPU 0.30-0.70, QPU >= 0.70)
# QPU_THRESHOLD: Entropy above which problem is "Quantum-Hard"
# GPU_THRESHOLD: Entropy above which problem is "GPU-Hard" (highly parallelizable)
QPU_THRESHOLD = GPU_MAX_SE if TENSOR_SE_AVAILABLE else 0.70
GPU_THRESHOLD = CPU_MAX_SE if TENSOR_SE_AVAILABLE else 0.30

# QUBO router imports (lazy loaded to avoid circular deps)
_qubo_router = None
_telemetry = None

def _get_qubo_router():
    """Lazy load QUBO router."""
    global _qubo_router
    if _qubo_router is None:
        from .qubo_router import qubo_router
        _qubo_router = qubo_router
    return _qubo_router

def _get_telemetry():
    """Lazy load telemetry observer."""
    global _telemetry
    if _telemetry is None:
        from .telemetry import telemetry
        _telemetry = telemetry
    return _telemetry


class SymbolicEntropyRouter:
    """
    Symbolic Entropy (Se) based routing logic for heterogeneous compute.

    In the Helix-CDC architecture, this component analyzes a BraidedTensor
    (or its underlying data/metadata) to calculate its complexity, Se,
    and dynamically routes the task to the optimal hardware (CPU, GPU, QPU).

    Se = H × U × D (real tensor metrics from tensor_classifier):
    - H = 1 - energy_at_10pct (entropy proxy)
    - U = 1 - neighbor_coherence (unstructuredness proxy, formerly "C")
    - D = sqrt(rank_ratio) (dimensionality proxy)

    2D Routing (NEW 2026-01-26):
    Uses (Se, C_struct) jointly for smarter backend selection:
    - Zone 1: Se < 0.30, C_struct >= 0.30 → CPU (simple + organized)
    - Zone 2: 0.30 <= Se < 0.70          → GPU (parallelizable)
    - Zone 3: Se >= 0.70, C_struct < 0.30 → QPU (complex + unstructured)
    - Zone 4: Se >= 0.70, C_struct >= 0.30 → GPU (complex but structured)

    Supports three modes:
    - Real Se mode (default): Uses tensor_se_estimator for H × U × D + 2D routing
    - Threshold mode (legacy): Fixed Se thresholds, variance-based Se
    - QUBO mode: Optimization considering Se + telemetry + thermal

    Set use_qubo_router=True to enable physics-aware routing.
    """

    def __init__(self, use_qubo_router: bool = False, gpu_vram_mb: float = 4000.0):
        """Initialize router with hardware availability status.

        Args:
            use_qubo_router: If True, use QUBO optimization instead of thresholds
            gpu_vram_mb: GPU VRAM in MB for size-based gating (default 4GB)
        """
        self.hardware_status = {
            'qpu_online': True,   # Assume QPU accessible for high Se tasks
            'gpu_online': True,   # Assume GPU accessible
            'cpu_online': True,   # CPU is always the fallback
        }

        # Configurable thresholds (used when use_qubo_router=False)
        self.qpu_threshold = QPU_THRESHOLD
        self.gpu_threshold = GPU_THRESHOLD

        # GPU VRAM limit for size-based gating
        self.gpu_vram_mb = gpu_vram_mb

        # QUBO routing mode
        self.use_qubo_router = use_qubo_router
        self._last_qubo_result = None  # For diagnostics

        # Track whether real Se is available
        self.use_real_se = TENSOR_SE_AVAILABLE
        self._last_se_result = None  # For diagnostics

    def _calculate_se(self, tensor_data: np.ndarray, state: dict) -> dict:
        """
        Calculate Symbolic Entropy (Se) and 2D routing info using REAL H × U × D.

        Se = H × U × D where:
        - H = 1 - energy_at_10pct (entropy: spread-out spectrum = high)
        - U = 1 - neighbor_coherence (unstructuredness: no spatial structure = high)
        - D = sqrt(rank_ratio) (depth: high effective rank = high)

        Returns dict with Se, C_struct, routing_zone, and routing_hint for 2D routing.
        Falls back to variance-based heuristic if tensor_se_estimator unavailable.
        """
        # Handle edge cases
        if tensor_data is None or tensor_data.size == 0:
            self._last_se_result = None
            return {'Se': 0.0, 'C_struct': 0.0, 'routing_zone': 1, 'routing_hint': 'cpu'}

        # Use REAL tensor Se if available (includes 2D routing)
        if self.use_real_se and TENSOR_SE_AVAILABLE:
            result = compute_routing_decision(
                tensor_data,
                state=state,
                gpu_vram_mb=self.gpu_vram_mb,
                name="tensor",
            )
            self._last_se_result = result
            return {
                'Se': result['Se'],
                'C_struct': result.get('C_struct', 0.0),
                'routing_zone': result.get('routing_zone', 2),
                'routing_hint': result.get('routing_hint', 'gpu'),
            }

        # FALLBACK: Variance-based heuristic (legacy behavior, no 2D routing)
        # 1. Base complexity from data variance (normalized to 0-0.5)
        std_dev = float(np.std(tensor_data))
        base_se = min(0.5, std_dev / 10.0)

        # 2. State-injected complexity (0-0.5)
        if state and isinstance(state, dict):
            problem_complexity = state.get('complexity_level', 0.0)
            state_se = min(0.5, problem_complexity / 10.0)
        else:
            state_se = 0.0

        # Total Se (normalized to 0-1)
        se = float(min(1.0, max(0.0, base_se + state_se)))

        # Fallback routing (1D, no C_struct)
        if se >= GPU_MAX_SE:
            hint = 'qpu'
            zone = 3
        elif se >= CPU_MAX_SE:
            hint = 'gpu'
            zone = 2
        else:
            hint = 'cpu'
            zone = 1

        self._last_se_result = {'Se': se, 'C_struct': 0.0, 'routing_zone': zone, 'fallback': True}
        return {'Se': se, 'C_struct': 0.0, 'routing_zone': zone, 'routing_hint': hint}

    def route_task(
        self,
        braided_tensor,
        state: dict = None,
        queue_depth: int = 0,
        deadline_ms: Optional[float] = None
    ) -> tuple:
        """
        Determine optimal hardware target based on tensor's Symbolic Entropy.

        2D Routing Logic (NEW 2026-01-26):
        Uses (Se, C_struct) jointly where C_struct = max(row_coherence, col_coherence):
            Zone 1: Se < 0.30, C_struct >= 0.30 → CPU (simple + organized)
            Zone 2: 0.30 <= Se < 0.70           → GPU (parallelizable)
            Zone 3: Se >= 0.70, C_struct < 0.30 → QPU (complex + unstructured)
            Zone 4: Se >= 0.70, C_struct >= 0.30 → GPU (complex but structured)

        Size Gate: If tensor > 80% GPU VRAM → CPU (won't fit)

        Args:
            braided_tensor: BraidedTensor object containing the data
            state: Optional dict with context/parameters (e.g., problem type)
            queue_depth: Number of pending tasks (used in QUBO mode)
            deadline_ms: Optional deadline pressure (used in QUBO mode)

        Returns:
            (backend: str, se_norm: float): Recommended backend and Se value
        """
        if state is None:
            state = {}

        # Get underlying numpy array for calculation
        try:
            tensor_data = braided_tensor.to_numpy()
        except Exception:
            # Fallback if conversion fails
            tensor_data = None

        # Calculate complexity metric (real Se = H × U × D) and 2D routing info
        se_result = self._calculate_se(tensor_data, state)
        se_norm = se_result['Se']
        c_struct = se_result.get('C_struct', 0.0)
        routing_hint = se_result.get('routing_hint', 'gpu')

        # --- SIZE GATE: Check if tensor fits in GPU VRAM ---
        size_gated = False
        if tensor_data is not None:
            size_mb = tensor_data.size * tensor_data.itemsize / 1e6
            if size_mb > self.gpu_vram_mb * 0.8:
                size_gated = True

        # --- QUBO MODE: Physics-aware optimization ---
        if self.use_qubo_router and not size_gated:
            try:
                qubo_router = _get_qubo_router()
                telemetry = _get_telemetry()

                # Get current physical state
                telem_snapshot = telemetry.snapshot()

                # Run QUBO optimization
                result = qubo_router.route(
                    se_norm=se_norm,
                    telemetry=telem_snapshot,
                    queue_depth=queue_depth,
                    deadline_ms=deadline_ms,
                    hardware_status=self.hardware_status
                )

                # Store for diagnostics
                self._last_qubo_result = result

                return result.backend, se_norm

            except ImportError as e:
                # Fallback to threshold mode if QUBO not available
                print(f"QUBO router unavailable, using 2D routing: {e}")

        # --- SIZE GATE OVERRIDE: Force CPU for large tensors ---
        if size_gated:
            return 'cpu', se_norm

        # --- 2D ROUTING: Use routing_hint from tensor_se_estimator ---
        # This already implements the (Se, C_struct) decision matrix
        if self.use_real_se and TENSOR_SE_AVAILABLE:
            # Respect hardware availability
            backend = routing_hint
            if backend == 'qpu' and not self.hardware_status['qpu_online']:
                backend = 'gpu'  # Fallback
            if backend == 'gpu' and not self.hardware_status['gpu_online']:
                backend = 'cpu'  # Fallback
            return backend, se_norm

        # --- LEGACY 1D Se-BASED ROUTING (fallback) ---
        if self.hardware_status['qpu_online'] and se_norm >= self.qpu_threshold:
            backend = 'qpu'
        elif self.hardware_status['gpu_online'] and se_norm >= self.gpu_threshold:
            backend = 'gpu'
        else:
            backend = 'cpu'

        return backend, se_norm

    def set_thresholds(self, gpu: float = None, qpu: float = None):
        """Update routing thresholds."""
        if gpu is not None:
            self.gpu_threshold = gpu
        if qpu is not None:
            self.qpu_threshold = qpu

    def set_hardware_status(self, cpu: bool = None, gpu: bool = None, qpu: bool = None):
        """Update hardware availability status."""
        if cpu is not None:
            self.hardware_status['cpu_online'] = cpu
        if gpu is not None:
            self.hardware_status['gpu_online'] = gpu
        if qpu is not None:
            self.hardware_status['qpu_online'] = qpu

    def info(self) -> dict:
        """Return router configuration and diagnostics."""
        mode = 'qubo' if self.use_qubo_router else ('real_se_2d' if self.use_real_se else 'threshold_1d')
        info = {
            'mode': mode,
            'use_real_se': self.use_real_se,
            'tensor_classifier_available': TENSOR_CLASSIFIER_AVAILABLE,
            'routing_type': '2D (Se, C_struct)' if self.use_real_se else '1D (Se only)',
            'thresholds': {
                'cpu_max_se': self.gpu_threshold,  # Se < this → CPU
                'gpu_max_se': self.qpu_threshold,  # Se < this → GPU, else QPU
                'struct_threshold': 0.30,  # C_struct >= this = structured
            },
            'zones': {
                1: 'CPU (Se<0.30)',
                2: 'GPU (0.30<=Se<0.70)',
                3: 'QPU (Se>=0.70, C_struct<0.30)',
                4: 'GPU (Se>=0.70, C_struct>=0.30, structured)',
            },
            'gpu_vram_mb': self.gpu_vram_mb,
            'hardware_status': self.hardware_status.copy(),
        }

        # Add last Se computation result
        if self._last_se_result:
            info['last_se_result'] = self._last_se_result

        # Add QUBO diagnostics if available
        if self.use_qubo_router and self._last_qubo_result:
            info['last_qubo_result'] = {
                'backend': self._last_qubo_result.backend,
                'energy': self._last_qubo_result.energy,
                'solve_time_ms': self._last_qubo_result.solve_time_ms,
                'all_energies': self._last_qubo_result.all_energies,
            }

        return info

    def get_last_qubo_result(self):
        """Return the last QUBO routing result for diagnostics."""
        return self._last_qubo_result

    def get_last_se_result(self):
        """Return the last Se computation result for diagnostics."""
        return self._last_se_result

    def enable_qubo_router(self, enable: bool = True):
        """Enable or disable QUBO routing mode."""
        self.use_qubo_router = enable

    def enable_real_se(self, enable: bool = True):
        """Enable or disable real Se computation (H × C × D).

        When disabled, falls back to variance-based heuristic.
        """
        self.use_real_se = enable and TENSOR_SE_AVAILABLE

    def set_gpu_vram(self, vram_mb: float):
        """Set GPU VRAM limit for size-based gating."""
        self.gpu_vram_mb = vram_mb


# Module-level singleton router instance (legacy threshold mode by default)
router = SymbolicEntropyRouter(use_qubo_router=False)

# QUBO-enabled router instance
qubo_enabled_router = SymbolicEntropyRouter(use_qubo_router=True)
