#!/usr/bin/env python3
"""
test_flowtorch_unit.py - Unit Tests for FlowTorch Routing

Tests the SymbolicEntropyRouter's CPU/GPU/QPU routing decisions.

Run:
    pytest tests/test_flowtorch_unit.py -v
"""

import sys
from pathlib import Path

import numpy as np
import pytest

# Add project root to path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from flowtorch.router import SymbolicEntropyRouter


class MockBraidedTensor:
    """Mock BraidedTensor for testing."""

    def __init__(self, data: np.ndarray):
        self._data = data

    def to_numpy(self) -> np.ndarray:
        return self._data


class TestSymbolicEntropyRouter:
    """Tests for SymbolicEntropyRouter."""

    def setup_method(self):
        """Set up router for each test."""
        self.router = SymbolicEntropyRouter(use_qubo_router=False)

    def test_router_creation(self):
        """Test router can be created."""
        router = SymbolicEntropyRouter()
        assert router is not None
        assert router.hardware_status['cpu_online']
        assert router.hardware_status['gpu_online']
        assert router.hardware_status['qpu_online']

    def test_router_info(self):
        """Test router info returns expected structure."""
        info = self.router.info()
        assert 'mode' in info
        assert 'thresholds' in info
        assert 'hardware_status' in info
        assert 'zones' in info

    def test_constant_tensor_routes_to_cpu(self):
        """Constant tensor should have low Se -> CPU."""
        # Constant tensor = low entropy = Se near 0
        data = np.ones((100, 100), dtype=np.float32)
        tensor = MockBraidedTensor(data)

        backend, se = self.router.route_task(tensor)

        # With variance-based fallback, constant -> low Se -> cpu
        assert backend == 'cpu', f"Expected cpu, got {backend}"
        assert se < 0.3, f"Expected low Se for constant, got {se}"

    def test_random_tensor_routes_higher(self):
        """Random tensor should have higher Se."""
        # Random tensor = high entropy
        np.random.seed(42)
        data = np.random.randn(100, 100).astype(np.float32) * 10

        tensor = MockBraidedTensor(data)
        backend, se = self.router.route_task(tensor)

        # Random should have higher Se than constant
        assert se > 0.0, f"Expected positive Se for random, got {se}"

    def test_hardware_offline_fallback(self):
        """Test fallback when hardware is offline."""
        # Set QPU and GPU offline
        self.router.set_hardware_status(qpu=False, gpu=False)

        # Even high-Se tensors should route to CPU
        np.random.seed(42)
        data = np.random.randn(100, 100).astype(np.float32) * 100

        tensor = MockBraidedTensor(data)
        backend, se = self.router.route_task(tensor)

        assert backend == 'cpu', f"Expected cpu fallback, got {backend}"

    def test_threshold_update(self):
        """Test threshold modification."""
        self.router.set_thresholds(gpu=0.2, qpu=0.5)

        info = self.router.info()
        assert info['thresholds']['cpu_max_se'] == 0.2
        assert info['thresholds']['gpu_max_se'] == 0.5

    def test_size_gating(self):
        """Test large tensor routes to CPU (VRAM gating)."""
        # Simulate tensor larger than GPU VRAM
        self.router.set_gpu_vram(1.0)  # 1 MB limit

        # Create tensor larger than 0.8 MB (80% of VRAM)
        data = np.random.randn(1000, 1000).astype(np.float32)
        tensor = MockBraidedTensor(data)

        backend, se = self.router.route_task(tensor)

        assert backend == 'cpu', "Large tensor should route to CPU"

    def test_empty_tensor_handled(self):
        """Test empty tensor doesn't crash."""
        data = np.array([], dtype=np.float32)
        tensor = MockBraidedTensor(data)

        backend, se = self.router.route_task(tensor)

        assert backend == 'cpu'
        assert se == 0.0

    def test_none_tensor_handled(self):
        """Test None data doesn't crash."""
        class NullTensor:
            def to_numpy(self):
                return None

        tensor = NullTensor()
        backend, se = self.router.route_task(tensor)

        assert backend == 'cpu'
        assert se == 0.0

    def test_state_complexity_injection(self):
        """Test state can inject complexity."""
        data = np.ones((100, 100), dtype=np.float32)
        tensor = MockBraidedTensor(data)

        # Without state: low Se
        backend1, se1 = self.router.route_task(tensor, state={})

        # With high complexity state: should increase Se (if variance fallback)
        backend2, se2 = self.router.route_task(
            tensor,
            state={'complexity_level': 20.0}
        )

        # Se should be higher with complexity injection
        assert se2 >= se1, f"Expected se2 >= se1, got {se2} < {se1}"


class TestRouterModes:
    """Test different router modes."""

    def test_qubo_mode_creation(self):
        """Test QUBO router mode can be enabled."""
        router = SymbolicEntropyRouter(use_qubo_router=True)
        info = router.info()
        assert info['mode'] == 'qubo'

    def test_mode_toggle(self):
        """Test mode can be toggled."""
        router = SymbolicEntropyRouter(use_qubo_router=False)
        assert router.info()['mode'] != 'qubo'

        router.enable_qubo_router(True)
        assert router.use_qubo_router is True

        router.enable_qubo_router(False)
        assert router.use_qubo_router is False

    def test_real_se_mode_toggle(self):
        """Test real Se computation mode toggle."""
        router = SymbolicEntropyRouter()

        # Disable real Se
        router.enable_real_se(False)
        assert router.use_real_se is False

        # Re-enable (if available)
        router.enable_real_se(True)
        # May or may not be available depending on imports


class TestRouterZones:
    """Test routing zone logic."""

    def test_zone_descriptions_exist(self):
        """Test zone descriptions are present."""
        router = SymbolicEntropyRouter()
        info = router.info()

        zones = info['zones']
        assert 1 in zones  # CPU zone
        assert 2 in zones  # GPU zone
        assert 3 in zones  # QPU zone
        assert 4 in zones  # GPU (structured) zone

    def test_last_se_result_available(self):
        """Test last Se result is stored for diagnostics."""
        router = SymbolicEntropyRouter()

        data = np.random.randn(50, 50).astype(np.float32)
        tensor = MockBraidedTensor(data)

        router.route_task(tensor)

        result = router.get_last_se_result()
        # May be None if real Se not available, that's OK
        # Just verify it doesn't crash


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
