"""Backend abstraction for TNFR mathematical kernels.

This module introduces a unified interface that maps core linear algebra
operations to concrete numerical libraries.  Keeping this layer small and
canonical guarantees we can switch implementations without diluting the
structural semantics required by TNFR (coherence, phase, νf, ΔNFR, etc.).

The canonical entry point is :func:`get_backend`, which honours three lookup
mechanisms in order of precedence:

1. Explicit ``name`` argument.
2. ``TNFR_MATH_BACKEND`` environment variable.
3. ``tnfr.backend_config.get_config().math_backend``.

The service configuration defaults to ``auto``. Automatic selection first
tries GPU-capable adapters in JAX → PyTorch → NumPy order, then repeats that
order for any available adapter, including CPU implementations. Request
``numpy`` explicitly to select it independently of installed optional adapters.
Factories initialize optional libraries only when a backend is requested.
A registered but unavailable named backend falls back to NumPy; an unknown
name raises ``LookupError``.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import (
    Any,
    Callable,
    ClassVar,
    Iterable,
    Mapping,
    MutableMapping,
    Protocol,
    cast,
    runtime_checkable,
)

from ..core.exceptions import BackendUnavailableError
from ..errors import TNFRValueError
from ..utils import cached_import, get_logger

logger = get_logger(__name__)


@runtime_checkable
class MathematicsBackend(Protocol):
    """Structural numerical backend interface.

    Notes
    -----
    Marked with @runtime_checkable to enable isinstance() checks for validating
    backend implementations conform to the expected mathematical operations interface.
    """

    name: str
    supports_autodiff: bool

    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any:
        """Convert ``value`` into a backend-native dense array."""

    def eig(self, matrix: Any) -> tuple[Any, Any]:
        """Return eigenvalues and eigenvectors for a general matrix."""

    def eigh(self, matrix: Any) -> tuple[Any, Any]:
        """Return eigenpairs for a Hermitian/symmetric matrix."""

    def matrix_exp(self, matrix: Any) -> Any:
        """Compute the matrix exponential of ``matrix``."""

    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any:
        """Return the matrix or vector norm according to ``ord``."""

    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any:
        """Evaluate an Einstein summation expression."""

    def matmul(self, a: Any, b: Any) -> Any:
        """Matrix multiplication that respects backend broadcasting rules."""

    def conjugate_transpose(self, matrix: Any) -> Any:
        """Hermitian conjugate of ``matrix`` († operator)."""

    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any:
        """Stack arrays along a new ``axis``."""

    def to_numpy(self, value: Any) -> Any:
        """Convert ``value`` to a ``numpy.ndarray`` when possible."""

    def is_gpu_available(self) -> bool:
        """Return True if the backend is currently using a GPU."""

    def get_device_name(self) -> str:
        """Return the name of the device being used (e.g. 'cpu', 'cuda:0')."""

    def get_backend_info(self) -> Mapping[str, Any]:
        """Return detailed backend information."""


BackendFactory = Callable[[], MathematicsBackend]


@dataclass(slots=True)
class _NumpyBackend:
    """NumPy backed implementation."""

    _np: Any
    _scipy_linalg: Any | None

    name: ClassVar[str] = "numpy"
    supports_autodiff: ClassVar[bool] = False

    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any:
        return self._np.asarray(value, dtype=dtype)

    def eig(self, matrix: Any) -> tuple[Any, Any]:
        return self._np.linalg.eig(matrix)

    def eigh(self, matrix: Any) -> tuple[Any, Any]:
        return self._np.linalg.eigh(matrix)

    def matrix_exp(self, matrix: Any) -> Any:
        """Evaluate general matrix exponentials through the required SciPy library."""
        if self._scipy_linalg is None:
            raise BackendUnavailableError(
                "SciPy is required for the NumPy matrix exponential"
            )
        return self._scipy_linalg.expm(matrix)

    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any:
        return self._np.linalg.norm(value, ord=ord, axis=axis)

    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any:
        return self._np.einsum(pattern, *operands, **kwargs)

    def matmul(self, a: Any, b: Any) -> Any:
        return self._np.matmul(a, b)

    def conjugate_transpose(self, matrix: Any) -> Any:
        return self._np.conjugate(matrix).T

    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any:
        return self._np.stack(tuple(arrays), axis=axis)

    def to_numpy(self, value: Any) -> Any:
        return self._np.asarray(value)

    def is_gpu_available(self) -> bool:
        return False

    def get_device_name(self) -> str:
        return "cpu"

    def get_backend_info(self) -> Mapping[str, Any]:
        return {
            "name": self.name,
            "version": self._np.__version__,
            "device": "cpu",
            "accelerated": False,
        }


@dataclass(slots=True)
class _JaxBackend:
    """JAX backed implementation."""

    _jnp: Any
    _jax_linalg: Any
    _jax: Any

    name: ClassVar[str] = "jax"
    supports_autodiff: ClassVar[bool] = True

    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any:
        return self._jnp.asarray(value, dtype=dtype)

    def eig(self, matrix: Any) -> tuple[Any, Any]:
        return self._jnp.linalg.eig(matrix)

    def eigh(self, matrix: Any) -> tuple[Any, Any]:
        return self._jnp.linalg.eigh(matrix)

    def matrix_exp(self, matrix: Any) -> Any:
        return self._jax_linalg.expm(matrix)

    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any:
        return self._jnp.linalg.norm(value, ord=ord, axis=axis)

    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any:
        return self._jnp.einsum(pattern, *operands, **kwargs)

    def matmul(self, a: Any, b: Any) -> Any:
        return self._jnp.matmul(a, b)

    def conjugate_transpose(self, matrix: Any) -> Any:
        return self._jnp.conjugate(matrix).T

    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any:
        return self._jnp.stack(tuple(arrays), axis=axis)

    def to_numpy(self, value: Any) -> Any:
        np_mod = cached_import("numpy")
        if np_mod is None:
            raise BackendUnavailableError("NumPy is required to export JAX arrays")
        return np_mod.asarray(self._jax.device_get(value))

    def is_gpu_available(self) -> bool:
        try:
            # Check if any device is a GPU or TPU
            return any(d.platform in ("gpu", "tpu") for d in self._jax.devices())
        except Exception:
            return False

    def get_device_name(self) -> str:
        try:
            return str(self._jax.devices()[0])
        except Exception:
            return "unknown"

    def get_backend_info(self) -> Mapping[str, Any]:
        return {
            "name": self.name,
            "version": self._jax.__version__,
            "device": self.get_device_name(),
            "accelerated": self.is_gpu_available(),
        }


@dataclass(slots=True)
class _TorchBackend:
    """PyTorch backed implementation with CUDA support."""

    _torch: Any
    _torch_linalg: Any
    _device: Any  # torch.device for CUDA/CPU placement
    _use_cuda: bool  # Whether CUDA is available and enabled

    name: ClassVar[str] = "torch"
    supports_autodiff: ClassVar[bool] = True

    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any:
        np_mod = cached_import("numpy")
        if (
            np_mod is not None
            and isinstance(value, np_mod.ndarray)
            and not value.flags.writeable
        ):
            # Torch tensors cannot preserve NumPy's read-only storage contract.
            value = value.copy(order="K")
        tensor = self._torch.as_tensor(value, device=self._device)
        if dtype is None:
            return tensor

        target_dtype = self._normalise_dtype(dtype)
        if target_dtype is None:
            return tensor.to(dtype=dtype, device=self._device)

        if tensor.dtype == target_dtype:
            return tensor.to(device=self._device)

        return tensor.to(dtype=target_dtype, device=self._device)

    def _normalise_dtype(self, dtype: Any) -> Any | None:
        """Return a ``torch.dtype`` equivalent for ``dtype`` when available."""

        if isinstance(dtype, self._torch.dtype):
            return dtype

        np_mod = cached_import("numpy")
        if np_mod is None:
            return None

        try:
            np_dtype = np_mod.dtype(dtype)
        except TypeError:
            return None

        numpy_name = np_dtype.name
        numpy_to_torch = {
            "bool": self._torch.bool,
            "uint8": self._torch.uint8,
            "int8": self._torch.int8,
            "int16": self._torch.int16,
            "int32": self._torch.int32,
            "int64": self._torch.int64,
            "float16": self._torch.float16,
            "float32": self._torch.float32,
            "float64": self._torch.float64,
            "complex64": getattr(self._torch, "complex64", None),
            "complex128": getattr(self._torch, "complex128", None),
            "bfloat16": getattr(self._torch, "bfloat16", None),
        }

        torch_dtype = numpy_to_torch.get(numpy_name)
        return torch_dtype

    def eig(self, matrix: Any) -> tuple[Any, Any]:
        eigenvalues, eigenvectors = self._torch.linalg.eig(matrix)
        return eigenvalues, eigenvectors

    def eigh(self, matrix: Any) -> tuple[Any, Any]:
        eigenvalues, eigenvectors = self._torch.linalg.eigh(matrix)
        return eigenvalues, eigenvectors

    def matrix_exp(self, matrix: Any) -> Any:
        return self._torch_linalg.matrix_exp(matrix)

    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any:
        if axis is None:
            return self._torch.linalg.norm(value, ord=ord)
        return self._torch.linalg.norm(value, ord=ord, dim=axis)

    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any:
        return self._torch.einsum(pattern, *operands, **kwargs)

    def matmul(self, a: Any, b: Any) -> Any:
        return self._torch.matmul(a, b)

    def conjugate_transpose(self, matrix: Any) -> Any:
        return matrix.mH if hasattr(matrix, "mH") else matrix.conj().transpose(-2, -1)

    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any:
        return self._torch.stack(tuple(arrays), dim=axis)

    def to_numpy(self, value: Any) -> Any:
        np_mod = cached_import("numpy")
        if np_mod is None:
            raise BackendUnavailableError("NumPy is required to export Torch tensors")
        if hasattr(value, "detach"):
            # Lazy conjugate/negative views cannot be exposed directly through
            # NumPy. Resolve their flags only on the detached observation;
            # native computations retain the original view and gradient graph.
            return value.detach().cpu().resolve_conj().resolve_neg().numpy()
        return np_mod.asarray(value)

    def is_gpu_available(self) -> bool:
        return self._use_cuda

    def get_device_name(self) -> str:
        return str(self._device)

    def get_backend_info(self) -> Mapping[str, Any]:
        return {
            "name": self.name,
            "version": self._torch.__version__,
            "device": self.get_device_name(),
            "accelerated": self.is_gpu_available(),
            "details": self.get_device_info(),
        }

    def get_device_info(self) -> dict[str, Any]:
        """Get CUDA device information."""
        info = {
            "device": str(self._device),
            "use_cuda": self._use_cuda,
            "cuda_available": (
                self._torch.cuda.is_available()
                if hasattr(self._torch, "cuda")
                else False
            ),
        }

        if self._use_cuda and hasattr(self._torch, "cuda"):
            info.update(
                {
                    "device_count": self._torch.cuda.device_count(),
                    "current_device": self._torch.cuda.current_device(),
                    "device_name": self._torch.cuda.get_device_name(
                        self._torch.cuda.current_device()
                    ),
                    "memory_allocated": self._torch.cuda.memory_allocated(),
                    "memory_reserved": self._torch.cuda.memory_reserved(),
                }
            )

        return info

    def to_cuda(self, value: Any) -> Any:
        """Move tensor to CUDA device if available."""
        if self._use_cuda and hasattr(value, "to"):
            return value.to(self._device)
        return value

    def to_cpu(self, value: Any) -> Any:
        """Move tensor to CPU device."""
        if hasattr(value, "cpu"):
            return value.cpu()
        return value


def _normalise_name(name: str) -> str:
    return name.strip().lower()


_BACKEND_FACTORIES: MutableMapping[str, BackendFactory] = {}
_BACKEND_ALIASES: MutableMapping[str, str] = {}
_BACKEND_CACHE: MutableMapping[str, MathematicsBackend] = {}

_AUTO_BACKEND_SENTINEL = "auto"
_AUTO_BACKEND_PRIORITY = ("jax", "torch", "numpy")


def _registration_name(name: str) -> str:
    """Normalize a name that explicit backend lookup can reach."""
    if not isinstance(name, str):
        raise TNFRValueError(
            "Backend identifiers must be nonempty strings other than 'auto'"
        )
    key = _normalise_name(name)
    if not key or key == _AUTO_BACKEND_SENTINEL:
        raise TNFRValueError(
            "Backend identifiers must be nonempty strings other than 'auto'"
        )
    return key


def ensure_array(
    value: Any,
    *,
    dtype: Any | None = None,
    backend: MathematicsBackend | None = None,
) -> Any:
    """Return ``value`` as a backend-native dense array."""

    resolved = backend or get_backend()
    return resolved.as_array(value, dtype=dtype)


def ensure_numpy(value: Any, *, backend: MathematicsBackend | None = None) -> Any:
    """Export ``value`` from the backend into :class:`numpy.ndarray`."""

    resolved = backend or get_backend()
    return resolved.to_numpy(value)


def register_backend(
    name: str,
    factory: BackendFactory,
    *,
    aliases: Iterable[str] | None = None,
    override: bool = False,
) -> None:
    """Register a backend factory under ``name``.

    Parameters
    ----------
    name:
        Nonempty canonical backend identifier, excluding the reserved ``auto``.
    factory:
        Callable that returns a :class:`MathematicsBackend` instance.
    aliases:
        Optional alternative identifiers that will resolve to ``name``.
        These have the same nonempty, non-reserved domain as canonical names.
    override:
        When ``True`` replaces an existing canonical factory or alias binding.
        Canonical names and aliases cannot shadow one another, even when
        overriding. Unmentioned aliases remain bound to their canonical name.

    Registration is validated before changing the registry. Replacing a
    canonical factory invalidates its cached instance, so its aliases also
    resolve to the replacement on their next lookup.
    """

    key = _registration_name(name)
    if key in _BACKEND_ALIASES:
        raise TNFRValueError(f"Backend name '{name}' is already an alias")
    if not override and key in _BACKEND_FACTORIES:
        raise TNFRValueError(
            f"Backend '{name}' already registered",
            context={"name": name, "existing": list(_BACKEND_FACTORIES.keys())},
            suggestion="Use override=True to replace the existing backend registration.",
        )
    proposed_aliases: dict[str, str] = {}
    if aliases is not None:
        for alias in aliases:
            alias_key = _registration_name(alias)
            if alias_key == key or alias_key in _BACKEND_FACTORIES:
                raise TNFRValueError(
                    f"Backend alias '{alias}' is already a canonical name"
                )
            if not override and (
                alias_key in _BACKEND_ALIASES or alias_key in proposed_aliases
            ):
                raise TNFRValueError(
                    f"Backend alias '{alias}' already registered",
                    context={
                        "alias": alias,
                        "existing_aliases": list(_BACKEND_ALIASES.keys()),
                    },
                    suggestion="Use override=True or choose a different alias.",
                )
            proposed_aliases[alias_key] = key

    _BACKEND_FACTORIES[key] = factory
    _BACKEND_ALIASES.update(proposed_aliases)
    _BACKEND_CACHE.pop(key, None)


def _resolve_backend_name(name: str | None) -> str:
    if name:
        normalised = _normalise_name(name)
        return (
            _AUTO_BACKEND_SENTINEL
            if normalised == _AUTO_BACKEND_SENTINEL
            else normalised
        )

    env_choice = os.getenv("TNFR_MATH_BACKEND")
    if env_choice:
        normalised = _normalise_name(env_choice)
        return (
            _AUTO_BACKEND_SENTINEL
            if normalised == _AUTO_BACKEND_SENTINEL
            else normalised
        )

    backend_from_config: str | None = None
    try:
        from ..backend_config import get_config  # Backend/GPU configuration

        config = get_config()
        backend_from_config = config.math_backend
    except Exception:  # pragma: no cover - defensive; config must not break selection
        backend_from_config = None

    if backend_from_config and backend_from_config != "auto":
        return _normalise_name(backend_from_config)

    return _AUTO_BACKEND_SENTINEL


def _resolve_factory(name: str) -> BackendFactory:
    canonical = _BACKEND_ALIASES.get(name, name)
    try:
        return _BACKEND_FACTORIES[canonical]
    except KeyError as exc:  # pragma: no cover - defensive path
        raise LookupError(f"Unknown mathematics backend: {name}") from exc


def _construct_backend(name: str) -> MathematicsBackend | None:
    canonical = _BACKEND_ALIASES.get(name, name)
    if canonical in _BACKEND_CACHE:
        return _BACKEND_CACHE[canonical]

    factory = _resolve_factory(canonical)
    try:
        backend = factory()
    except BackendUnavailableError as exc:
        logger.warning("Backend '%s' unavailable: %s", canonical, exc)
        return None

    _BACKEND_CACHE[canonical] = backend
    return backend


def get_backend(name: str | None = None) -> MathematicsBackend:
    """Return a backend instance using the configured resolution order."""

    resolved_name = _resolve_backend_name(name)
    if resolved_name == _AUTO_BACKEND_SENTINEL:
        # First pass: Look for GPU-accelerated backend
        for candidate in _AUTO_BACKEND_PRIORITY:
            backend = _construct_backend(candidate)
            if backend is not None and backend.is_gpu_available():
                logger.info(
                    "Auto-selected GPU-accelerated backend '%s' (%s)",
                    candidate,
                    backend.get_device_name(),
                )
                return backend

        # Second pass: Fallback to any available backend
        for candidate in _AUTO_BACKEND_PRIORITY:
            backend = _construct_backend(candidate)
            if backend is not None:
                if candidate != "numpy":
                    logger.info("Auto-selected CPU backend '%s'", candidate)
                return backend
        raise BackendUnavailableError(
            "No mathematical backend available; tried: "
            + ", ".join(_AUTO_BACKEND_PRIORITY)
        )

    backend = _construct_backend(resolved_name)
    if backend is not None:
        return backend

    if resolved_name != "numpy":
        logger.warning("Falling back to NumPy backend")
        fallback = _construct_backend("numpy")
        if fallback is not None:
            return fallback

    raise BackendUnavailableError(
        f"Unable to initialise mathematics backend '{resolved_name}'"
    )


def available_backends() -> Mapping[str, BackendFactory]:
    """Return the registered backend factories."""

    return dict(_BACKEND_FACTORIES)


def _make_numpy_backend() -> MathematicsBackend:
    np_module = cached_import("numpy")
    if np_module is None:
        raise BackendUnavailableError("NumPy is not installed")
    scipy_linalg = cached_import("scipy.linalg")
    if scipy_linalg is None:
        logger.debug("SciPy not available; NumPy matrix exponential is unavailable")
    backend = _NumpyBackend(np_module, scipy_linalg)  # type: ignore[call-arg]
    return cast(MathematicsBackend, backend)


def _make_jax_backend() -> MathematicsBackend:
    jnp_module = cached_import("jax.numpy")
    if jnp_module is None:
        raise BackendUnavailableError("jax.numpy is not available")
    jax_scipy = cached_import("jax.scipy.linalg")
    if jax_scipy is None:
        raise BackendUnavailableError("jax.scipy.linalg is required for matrix_exp")
    jax_module = cached_import("jax")
    if jax_module is None:
        raise BackendUnavailableError("jax core module is required")
    # Enable 64-bit precision: TNFR is canonical float64 numerical physics
    # (mpmath-derived constants, conservation residuals, eigenvalue analysis).
    # Without this, jnp.asarray(x, dtype=float64) silently downcasts to float32
    # and emits a UserWarning, breaking precision contracts of structural
    # field telemetry (Phi_s, |grad phi|, K_phi, xi_C). Idempotent call.
    try:
        jax_config = getattr(jax_module, "config", None)
        if jax_config is not None and not bool(
            getattr(jax_config, "jax_enable_x64", False)
        ):
            jax_config.update("jax_enable_x64", True)
    except Exception:
        # Non-fatal: if x64 cannot be enabled the backend still works
        # in 32-bit (with the existing dtype downcast warnings).
        pass
    backend = _JaxBackend(jnp_module, jax_scipy, jax_module)  # type: ignore[call-arg]
    return cast(MathematicsBackend, backend)


def _make_torch_backend() -> MathematicsBackend:
    torch_module = cached_import("torch")
    if torch_module is None:
        raise BackendUnavailableError("PyTorch is not installed")
    torch_linalg = cached_import("torch.linalg")
    if torch_linalg is None:
        raise BackendUnavailableError(
            "torch.linalg is required for linear algebra operations"
        )

    # CUDA device detection and configuration
    use_cuda = False
    device = torch_module.device("cpu")

    if hasattr(torch_module, "cuda") and torch_module.cuda.is_available():
        # Check environment variable for CUDA preference
        cuda_enabled = os.getenv("TNFR_CUDA_ENABLED", "true").lower() in (
            "true",
            "1",
            "yes",
        )
        if cuda_enabled:
            try:
                device = torch_module.device("cuda")
                use_cuda = True
                logger.info(
                    f"CUDA enabled: {torch_module.cuda.get_device_name()} with {torch_module.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB"
                )
            except Exception as e:
                logger.warning(
                    f"CUDA available but failed to initialize: {e}. Falling back to CPU."
                )
                device = torch_module.device("cpu")
        else:
            logger.info("CUDA available but disabled via TNFR_CUDA_ENABLED=false")
    else:
        logger.info("CUDA not available, using CPU backend")

    backend = _TorchBackend(torch_module, torch_linalg, device, use_cuda)  # type: ignore[call-arg]
    return cast(MathematicsBackend, backend)


register_backend("numpy", _make_numpy_backend, aliases=("np",))
register_backend("jax", _make_jax_backend)
register_backend("torch", _make_torch_backend, aliases=("pytorch",))

__all__ = [
    "MathematicsBackend",
    "BackendUnavailableError",
    "register_backend",
    "get_backend",
    "available_backends",
    "ensure_array",
    "ensure_numpy",
]