#!/usr/bin/env python3
"""
FlowTorch Braid - Unified Multi-Framework Tensor Bridge (DLPack Zero-Copy)

This module implements the FlowBraid class and BraidedTensor, which allows
seamless, zero-copy data transfer between PyTorch, TensorFlow, JAX, CuPy, and NumPy.
It acts as the Universal Driver Runner (UDR) interface for the Symbolic Entropy
(Se) based scheduler, dynamically routing computation based on tensor complexity.

The core technology relies on the DLPack standard (https://github.com/dmlc/dlpack).

Usage:
    from flowtorch.braid import FlowBraid, get_braid

    braid = get_braid()

    # Ingest ANY tensor from ANY framework
    unified = braid.ingest(pytorch_tensor)
    unified = braid.ingest(tensorflow_tensor)
    unified = braid.ingest(jax_array)

    # Export to ANY framework (zero-copy via DLPack)
    pytorch_out = braid.to_torch(unified)
    tf_out = braid.to_tf(unified)

    # Route through Se-based scheduler
    output, backend, se = braid.route(unified, {'complexity_level': 5})
"""

import sys
import os
import numpy as np

# Attempt graceful imports for all supported frameworks
# We use try-except blocks with multiple exception types (ImportError, AttributeError, etc.)
# to handle issues like TensorFlow/JAX crashing due to incompatible NumPy 2.x versions.

_HAS_TORCH = False
try:
    import torch
    from torch.utils import dlpack as tdl
    _HAS_TORCH = True
except ImportError:
    torch = None
    tdl = None

_HAS_TF = False
_TF_IMPORT_ERROR = None
try:
    # Suppress TensorFlow warnings
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    import tensorflow as tf
    _HAS_TF = True
except (ImportError, AttributeError, TypeError, Exception) as e:
    # TensorFlow can crash on NumPy 2.x incompatibility with AttributeError/TypeError
    tf = None
    _TF_IMPORT_ERROR = str(e) if str(e) else "TensorFlow import failed"

_HAS_JAX = False
_JAX_IMPORT_ERROR = None
try:
    import jax
    import jax.numpy as jnp
    import jax.dlpack as jdl
    _HAS_JAX = True
except (ImportError, AttributeError, TypeError, Exception) as e:
    jax = None
    jnp = None
    jdl = None
    _JAX_IMPORT_ERROR = str(e) if str(e) else "JAX import failed"

_HAS_CUPY = False
try:
    import cupy as cp
    _HAS_CUPY = True
except ImportError:
    cp = None

# Import the Se router
_HAS_ROUTER = False
try:
    from flowtorch.router import router, SymbolicEntropyRouter
    _HAS_ROUTER = True
except ImportError:
    # Create a mock router if the file is missing
    class MockRouter:
        def route_task(self, *args, **kwargs):
            return 'cpu', 0.0
        def info(self):
            return {'thresholds': {'gpu': 0.45, 'qpu': 0.85}, 'hardware_status': {}}
    router = MockRouter()


class Framework:
    """Enum-like class for framework identifiers."""
    TORCH = 'torch'
    TENSORFLOW = 'tf'
    JAX = 'jax'
    CUPY = 'cupy'
    NUMPY = 'numpy'


class BraidedTensor:
    """
    A unified tensor representation that holds data and metadata
    for zero-copy transfer via DLPack, acting as the UDR data carrier.

    Note: DLPack capsules can only be consumed once. We keep the original
    tensor reference to allow multiple conversions.
    """

    def __init__(self, tensor, original_framework):
        self.original_framework = original_framework
        self.shape = tuple(tensor.shape)
        self.dtype = str(tensor.dtype)

        # Store the original tensor for reuse (DLPack capsules are single-use)
        self._original_tensor = tensor
        self._original_numpy = None  # Fallback storage for when no DLPack available

    def to_torch(self):
        """Convert to PyTorch tensor."""
        if not _HAS_TORCH:
            raise RuntimeError("PyTorch not available")

        src = self.original_framework
        tensor = self._original_tensor

        if src == Framework.TORCH:
            return tensor
        elif src == Framework.NUMPY:
            return torch.from_numpy(tensor)
        elif src == Framework.CUPY and _HAS_CUPY:
            return torch.as_tensor(tensor)
        else:
            # Fallback via numpy
            return torch.from_numpy(self.to_numpy())

    def to_tf(self):
        """Convert to TensorFlow tensor."""
        if not _HAS_TF:
            raise RuntimeError("TensorFlow not available")

        src = self.original_framework
        tensor = self._original_tensor

        if src == Framework.TENSORFLOW:
            return tensor
        elif src == Framework.NUMPY:
            return tf.constant(tensor)
        elif src == Framework.TORCH and _HAS_TORCH:
            # Torch → TF via DLPack (fresh capsule)
            dl = tdl.to_dlpack(tensor.contiguous())
            return tf.experimental.dlpack.from_dlpack(dl)
        else:
            return tf.constant(self.to_numpy())

    def to_jax(self):
        """Convert to JAX array."""
        if not _HAS_JAX:
            raise RuntimeError("JAX not available")

        src = self.original_framework
        tensor = self._original_tensor

        if src == Framework.JAX:
            return tensor
        elif src == Framework.NUMPY:
            return jnp.array(tensor)
        elif src == Framework.TORCH and _HAS_TORCH:
            dl = tdl.to_dlpack(tensor.contiguous())
            return jdl.from_dlpack(dl)
        else:
            return jnp.array(self.to_numpy())

    def to_cupy(self):
        """Convert to CuPy array."""
        if not _HAS_CUPY:
            raise RuntimeError("CuPy not available")

        src = self.original_framework
        tensor = self._original_tensor

        if src == Framework.CUPY:
            return tensor
        elif src == Framework.NUMPY:
            return cp.array(tensor)
        elif src == Framework.TORCH and _HAS_TORCH:
            return cp.asarray(tensor)
        else:
            return cp.array(self.to_numpy())

    def to_numpy(self):
        """Convert to NumPy array (may require copy from GPU)."""
        src = self.original_framework
        tensor = self._original_tensor

        if src == Framework.NUMPY:
            return tensor
        elif src == Framework.TORCH and _HAS_TORCH:
            return tensor.detach().cpu().numpy()
        elif src == Framework.TENSORFLOW and _HAS_TF:
            return tensor.numpy()
        elif src == Framework.JAX and _HAS_JAX:
            return np.array(tensor)
        elif src == Framework.CUPY and _HAS_CUPY:
            return tensor.get()
        else:
            raise ValueError(f"Cannot convert {src} to NumPy")


class FlowBraid:
    """
    The main FlowTorch interface. Manages tensor ingestion, export, and routing.
    """

    def __init__(self):
        self.available_frameworks = self._check_frameworks()
        if not _HAS_ROUTER:
            print("[Warning] SymbolicEntropyRouter not found. Routing will default to CPU.")

    def _check_frameworks(self):
        """Returns a dict of available frameworks and their status/error."""
        return {
            'torch_available': _HAS_TORCH,
            'tensorflow_available': _HAS_TF,
            'tensorflow_error': _TF_IMPORT_ERROR if not _HAS_TF and _TF_IMPORT_ERROR else None,
            'jax_available': _HAS_JAX,
            'jax_error': _JAX_IMPORT_ERROR if not _HAS_JAX and _JAX_IMPORT_ERROR else None,
            'cupy_available': _HAS_CUPY,
            'numpy_available': True,
            'router_available': _HAS_ROUTER,
        }

    def info(self):
        """Return status of all framework integrations and the router."""
        result = {
            'frameworks': {},
            'router': router.info() if _HAS_ROUTER else {'status': 'mock'},
        }
        for key, value in self.available_frameworks.items():
            if 'error' not in key:
                result['frameworks'][key] = value
        return result

    def print_info(self):
        """Print the status of all framework integrations."""
        print("--- FlowBraid System Status ---")
        for key, value in self.available_frameworks.items():
            if 'error' in key and value:
                print(f"  {key.replace('_error', '')}: Failed ({value[:50]}...)")
            elif 'available' in key:
                print(f"  {key}: {value}")
        print(f"  router_available: {_HAS_ROUTER}")
        print("-------------------------------")

    def ingest(self, tensor) -> BraidedTensor:
        """
        Ingest a tensor from any supported framework and convert to BraidedTensor.

        Args:
            tensor: Tensor from PyTorch, TensorFlow, JAX, CuPy, or NumPy

        Returns:
            BraidedTensor for zero-copy transfers and routing
        """
        if _HAS_TORCH and isinstance(tensor, torch.Tensor):
            framework = Framework.TORCH
        elif _HAS_TF and tf is not None and hasattr(tf, 'is_tensor') and tf.is_tensor(tensor):
            framework = Framework.TENSORFLOW
        elif _HAS_JAX and jnp is not None and hasattr(jax, 'Array') and isinstance(tensor, jax.Array):
            framework = Framework.JAX
        elif _HAS_CUPY and isinstance(tensor, cp.ndarray):
            framework = Framework.CUPY
        elif isinstance(tensor, np.ndarray):
            framework = Framework.NUMPY
        else:
            raise TypeError(f"Unsupported input tensor type: {type(tensor)}")

        return BraidedTensor(tensor, framework)

    def route(self, braided_tensor: BraidedTensor, state: dict = None) -> tuple:
        """
        Route task using Symbolic Entropy (Se) to determine optimal backend.

        The Se metric determines the optimal hardware backend:
        - If Se >= QPU_THRESHOLD (0.85): Route to Quantum Processing Unit
        - If Se >= GPU_THRESHOLD (0.45): Route to Graphics Processing Unit
        - Else: Route to CPU

        Args:
            braided_tensor: BraidedTensor to route
            state: Dict with context (e.g., {'complexity_level': 0-10})

        Returns:
            (output_tensor, backend: str, se_norm: float)
        """
        if state is None:
            state = {}

        # Determine target backend and Se complexity score
        target_backend, se_norm = router.route_task(braided_tensor, state)

        # Convert BraidedTensor to format required by target backend
        output_tensor = None

        try:
            if target_backend == 'qpu':
                # QPU uses PyTorch or JAX as input to quantum circuit compiler
                if _HAS_TORCH:
                    output_tensor = braided_tensor.to_torch()
                elif _HAS_JAX:
                    output_tensor = braided_tensor.to_jax()
                else:
                    output_tensor = braided_tensor.to_numpy()
                    target_backend = f"qpu_via_{Framework.NUMPY}"

            elif target_backend == 'gpu':
                try:
                    if _HAS_CUPY:
                        output_tensor = braided_tensor.to_cupy()
                    elif _HAS_TORCH and torch.cuda.is_available():
                        output_tensor = braided_tensor.to_torch().cuda()
                    elif _HAS_TF:
                        output_tensor = braided_tensor.to_tf()
                    else:
                        output_tensor = braided_tensor.to_numpy()
                        target_backend = f"gpu_via_{Framework.NUMPY}"
                except Exception:
                    # GPU failed, fall back to CPU
                    output_tensor = braided_tensor.to_numpy()
                    target_backend = "cpu_fallback"

            elif target_backend == 'cpu':
                output_tensor = braided_tensor.to_numpy()

            else:
                output_tensor = braided_tensor.to_numpy()
                target_backend = 'cpu'

        except Exception:
            # Ultimate fallback
            output_tensor = braided_tensor.to_numpy()
            target_backend = 'cpu_fallback'

        return output_tensor, target_backend, se_norm

    # --- Convenience Export Methods ---

    def to_torch(self, braided_tensor: BraidedTensor):
        """Convert BraidedTensor to PyTorch."""
        return braided_tensor.to_torch()

    def to_tf(self, braided_tensor: BraidedTensor):
        """Convert BraidedTensor to TensorFlow."""
        return braided_tensor.to_tf()

    def to_jax(self, braided_tensor: BraidedTensor):
        """Convert BraidedTensor to JAX."""
        return braided_tensor.to_jax()

    def to_cupy(self, braided_tensor: BraidedTensor):
        """Convert BraidedTensor to CuPy."""
        return braided_tensor.to_cupy()

    def to_numpy(self, braided_tensor: BraidedTensor):
        """Convert BraidedTensor to NumPy."""
        return braided_tensor.to_numpy()


# Module-level singleton
_default_braid = None

def get_braid() -> FlowBraid:
    """Returns the FlowBraid singleton."""
    global _default_braid
    if _default_braid is None:
        _default_braid = FlowBraid()
    return _default_braid


# --- Example Usage / Verification Script ---

if __name__ == '__main__':
    print("=== FlowBraid + Symbolic Entropy Router Test ===\n")

    braid = FlowBraid()
    braid.print_info()

    # Define states with different complexity levels
    simple_state = {'task_name': 'Vector Add', 'complexity_level': 1.0}
    parallel_state = {'task_name': 'Matrix Multiply (10k)', 'complexity_level': 5.0}
    quantum_state = {'task_name': 'QUBO Optimization', 'complexity_level': 9.5}

    # ------------------------------------------------------------------
    # TEST 1: Simple Task (CPU Route Expected)
    print("\n[TEST 1] Simple CPU Task")
    data_simple = np.array([i * 0.1 for i in range(10)])  # Low variance
    unified_simple = braid.ingest(data_simple)
    output_simple, backend_simple, se_simple = braid.route(unified_simple, simple_state)

    assert backend_simple == 'cpu', f"Expected 'cpu', got '{backend_simple}'"
    print(f"✅ Route 1 Succeeded: {backend_simple} (Se: {se_simple:.4f})")
    print(f"   Output Type: {type(output_simple)}")

    # ------------------------------------------------------------------
    # TEST 2: High Parallel Task (GPU Route Expected)
    print("\n[TEST 2] High Parallel GPU Task")
    data_parallel = np.random.rand(100, 100) * 5.0  # Higher variance
    unified_parallel = braid.ingest(data_parallel)
    output_parallel, backend_parallel, se_parallel = braid.route(unified_parallel, parallel_state)

    if 'gpu' in backend_parallel:
        print(f"✅ Route 2 Succeeded: {backend_parallel} (Se: {se_parallel:.4f})")
    else:
        print(f"⚠️ Route 2 Succeeded (Fallback): {backend_parallel} (Se: {se_parallel:.4f})")
    print(f"   Output Type: {type(output_parallel)}")

    # ------------------------------------------------------------------
    # TEST 3: Quantum Hard Task (QPU Route Expected)
    print("\n[TEST 3] Quantum-Hard QPU Task")
    data_quantum = np.random.normal(loc=10.0, scale=20.0, size=(5, 5))
    unified_quantum = braid.ingest(data_quantum)
    output_quantum, backend_quantum, se_quantum = braid.route(unified_quantum, quantum_state)

    assert 'qpu' in backend_quantum, f"Expected 'qpu', got '{backend_quantum}'"
    print(f"✅ Route 3 Succeeded: {backend_quantum} (Se: {se_quantum:.4f})")
    print(f"   Output Type: {type(output_quantum)}")

    # ------------------------------------------------------------------
    # TEST 4: Integrity Check (data preserved through routing)
    if _HAS_TORCH:
        print("\n[TEST 4] Data Integrity Check")
        torch_tensor = torch.randn(2, 2)
        original_np = torch_tensor.numpy()

        braid_test = braid.ingest(torch_tensor)
        output_routed, backend_test, se_test = braid.route(braid_test, quantum_state)

        # Convert output back to numpy for comparison
        if isinstance(output_routed, torch.Tensor):
            output_np = output_routed.numpy()
        else:
            output_np = output_routed

        # Check data integrity
        max_diff = np.abs(original_np - output_np).max()
        assert max_diff < 1e-6, f"Data corrupted: max diff = {max_diff}"
        print(f"✅ Data Integrity Check Passed. Backend: {backend_test} (max_diff: {max_diff:.2e})")

    print("\n=== All Tests Passed ===")
