dynamics.py

"""Spectral dynamics helpers driven by ΔNFR generators."""

from __future__ import annotations

from dataclasses import field
from typing import Any, NamedTuple, Sequence

from ..compat.dataclass import dataclass
from .backend import MathematicsBackend, ensure_array, ensure_numpy, get_backend
from .spaces import HilbertSpace
from .unified_numerical import TNFRValueError, np

try:  # pragma: no cover - optional SciPy dependency
    from scipy.linalg import expm as _scipy_expm  # type: ignore
except Exception:  # pragma: no cover - SciPy not installed
    _scipy_expm = None

__all__ = ["MathematicalDynamicsEngine", "ContractiveDynamicsEngine"]


def _has_backend_matrix_exp(backend: MathematicsBackend) -> bool:
    """Return ``True`` when ``backend`` exposes a usable ``matrix_exp``."""

    matrix_exp = getattr(backend, "matrix_exp", None)
    if not callable(matrix_exp):
        return False

    try:
        probe = ensure_array([[0.0]], dtype=np.complex128, backend=backend)
        matrix_exp(probe)
    except (AttributeError, NotImplementedError):
        return False
    except Exception:
        # Older backends may surface missing implementations as runtime errors;
        # treat them as signals to fall back to SciPy when available.
        return False
    return True


def _as_matrix(
    matrix: Sequence[Sequence[complex]] | np.ndarray | Any,
    *,
    backend: MathematicsBackend,
) -> Any:
    arr = ensure_array(matrix, dtype=np.complex128, backend=backend)
    shape = getattr(arr, "shape", None)
    if shape is None or len(shape) != 2 or shape[0] != shape[1]:
        raise TNFRValueError(
            "Generator matrix must be square.",
            context={"shape": shape},
            suggestion="Provide a square matrix.",
        )
    return arr


def _is_hermitian(
    matrix: Any, *, atol: float = 1e-9, backend: MathematicsBackend
) -> bool:
    matrix_np = ensure_numpy(matrix, backend=backend)
    return bool(np.allclose(matrix_np, matrix_np.conj().T, atol=atol))


def _vectorize_density(matrix: Any, *, backend: MathematicsBackend) -> Any:
    arr = ensure_array(matrix, dtype=np.complex128, backend=backend)
    return arr.transpose(1, 0).reshape((-1,))


def _devectorize_density(vector: Any, dim: int, *, backend: MathematicsBackend) -> Any:
    arr = ensure_array(vector, dtype=np.complex128, backend=backend)
    return arr.reshape((dim, dim)).transpose(1, 0)


class TraceValue(NamedTuple):
    """Container for trace evaluations in both backend and NumPy space."""

    backend: Any
    numpy: complex | None


def _trace(matrix: Any, *, backend: MathematicsBackend) -> TraceValue:
    traced_backend = backend.einsum("ii->", matrix)
    try:
        traced_numpy = complex(
            np.asarray(ensure_numpy(traced_backend, backend=backend))
        )
    except (ValueError, TypeError, Exception):
        # Fallback for backends where conversion fails (e.g. JAX tracing)
        traced_numpy = None
    return TraceValue(traced_backend, traced_numpy)


@dataclass(slots=True)
class MathematicalDynamicsEngine:
    """Unitary evolution generated by Hermitian ΔNFR operators.

    The engine accepts inputs expressed as backend-native tensors (NumPy,
    :mod:`jax`, :mod:`torch`).  When the configured backend supports automatic
    differentiation the evolution map ``exp(-i·Δ·dt)`` remains differentiable
    because native propagators are now preferred.  Passing ``use_scipy=True``
    explicitly opts into SciPy's exponential; we only fall back automatically
    when the backend lacks a ``matrix_exp`` implementation.
    """

    generator: np.ndarray
    hilbert_space: HilbertSpace
    atol: float = 1e-9
    _use_scipy: bool = False
    backend: MathematicsBackend = field(init=False, repr=False)
    _generator_backend: Any = field(init=False, repr=False)
    _numpy_generator: np.ndarray = field(init=False, repr=False)

    def __init__(
        self,
        generator: Sequence[Sequence[complex]] | np.ndarray | Any,
        hilbert_space: HilbertSpace,
        *,
        atol: float = 1e-9,
        use_scipy: bool | None = None,
        backend: MathematicsBackend | None = None,
    ) -> None:
        resolved_backend = backend or get_backend()
        matrix = _as_matrix(generator, backend=resolved_backend)
        matrix_np = ensure_numpy(matrix, backend=resolved_backend)
        if matrix_np.shape != (hilbert_space.dimension, hilbert_space.dimension):
            raise TNFRValueError(
                "Generator dimension must match the Hilbert space.",
                context={
                    "generator_shape": matrix_np.shape,
                    "hilbert_dimension": hilbert_space.dimension,
                },
                suggestion="Ensure generator matches Hilbert space dimension.",
            )
        if not _is_hermitian(matrix, atol=atol, backend=resolved_backend):
            raise TNFRValueError(
                "Dynamics generator must be Hermitian.",
                context={"atol": atol},
                suggestion="Ensure generator is Hermitian.",
            )
        self.backend = resolved_backend
        self._generator_backend = matrix
        self._numpy_generator = matrix_np
        self.generator = matrix_np
        self.hilbert_space = hilbert_space
        self.atol = float(atol)
        if use_scipy is None:
            has_matrix_exp = _has_backend_matrix_exp(self.backend)
            if has_matrix_exp:
                self._use_scipy = False
            elif _scipy_expm is not None:
                self._use_scipy = True
            else:
                raise RuntimeError(
                    "Backend lacks matrix_exp and SciPy is unavailable for fallback."
                )
        else:
            if use_scipy and _scipy_expm is None:
                raise RuntimeError("SciPy expm requested but SciPy is not available.")
            self._use_scipy = bool(use_scipy and _scipy_expm is not None)

    def _unitary_backend(self, dt: float) -> Any:
        if self._use_scipy and _scipy_expm is not None:
            return ensure_array(
                _scipy_expm(-1j * dt * self._numpy_generator),
                backend=self.backend,
            )
        return self.backend.matrix_exp(-1j * dt * self._generator_backend)

    def step(
        self,
        state: Sequence[complex] | np.ndarray | Any,
        *,
        dt: float = 1.0,
        normalize: bool = True,
    ) -> Any:
        """Evolve ``state`` by ``dt`` using the unitary ``exp(-i·Δ·dt)``."""

        vector = ensure_array(state, dtype=np.complex128, backend=self.backend)
        if vector.shape != (self.hilbert_space.dimension,):
            raise TNFRValueError(
                "State vector dimension mismatch.",
                context={
                    "vector_shape": vector.shape,
                    "expected_dimension": self.hilbert_space.dimension,
                },
                suggestion="Ensure state vector matches Hilbert space dimension.",
            )
        unitary = self._unitary_backend(dt)
        evolved = self.backend.matmul(unitary, vector)
        if normalize:
            norm_backend = self.backend.norm(evolved)
            norm_numpy = float(
                np.asarray(ensure_numpy(norm_backend, backend=self.backend))
            )
            if np.isclose(norm_numpy, 0.0, atol=self.atol):
                raise TNFRValueError(
                    "Cannot normalise a null state vector.",
                    context={"norm": norm_numpy, "atol": self.atol},
                    suggestion="Provide a non-zero state vector.",
                )
            evolved = evolved / norm_backend
        return evolved

    def evolve(
        self,
        state: Sequence[complex] | np.ndarray | Any,
        *,
        steps: int,
        dt: float = 1.0,
        normalize: bool = True,
    ) -> Any:
        """Return trajectory of length ``steps + 1`` starting from ``state``."""

        if steps < 0:
            raise TNFRValueError(
                "steps must be non-negative.",
                context={"steps": steps},
                suggestion="Provide a non-negative integer for steps.",
            )
        current = ensure_array(state, dtype=np.complex128, backend=self.backend)
        if current.shape != (self.hilbert_space.dimension,):
            raise TNFRValueError(
                "State dimension mismatch.",
                context={
                    "expected_dimension": self.hilbert_space.dimension,
                    "received_shape": current.shape,
                },
                suggestion="Ensure state vector matches Hilbert space dimension.",
            )
        trajectory: list[Any] = [current]
        for _ in range(steps):
            current = self.step(current, dt=dt, normalize=normalize)
            trajectory.append(current)
        return self.backend.stack(trajectory, axis=0)


@dataclass(slots=True)
class ContractiveDynamicsEngine:
    """Contractive semigroup evolution driven by Lindblad ΔNFR generators.

    Backend-native tensors are accepted for all density operators.  When the
    chosen backend supports automatic differentiation we keep gradients intact
    by default because native semigroup propagators are preferred.  Requesting
    ``use_scipy=True`` still falls back to SciPy's :func:`scipy.linalg.expm`,
    primarily for generators missing backend support.
    """

    generator: np.ndarray
    hilbert_space: HilbertSpace
    atol: float = 1e-9
    _use_scipy: bool = False
    backend: MathematicsBackend = field(init=False, repr=False)
    _generator_backend: Any = field(init=False, repr=False)
    _numpy_generator: np.ndarray = field(init=False, repr=False)
    _identity_backend: Any = field(init=False, repr=False)
    _last_contractivity_gap: float = field(init=False, repr=False)

    def __init__(
        self,
        generator: Sequence[Sequence[complex]] | np.ndarray | Any,
        hilbert_space: HilbertSpace,
        *,
        atol: float = 1e-9,
        ensure_contractive: bool = True,
        use_scipy: bool | None = None,
        backend: MathematicsBackend | None = None,
    ) -> None:
        resolved_backend = backend or get_backend()
        matrix = _as_matrix(generator, backend=resolved_backend)
        matrix_np = ensure_numpy(matrix, backend=resolved_backend)
        expected = hilbert_space.dimension * hilbert_space.dimension
        if matrix_np.shape != (expected, expected):
            raise TNFRValueError(
                "Generator must act on vectorised density operators.",
                context={
                    "expected_dimension": (expected, expected),
                    "received_shape": matrix_np.shape,
                },
                suggestion="Ensure generator dimension matches vectorized Hilbert space.",
            )
        self.backend = resolved_backend
        self._generator_backend = matrix
        self._numpy_generator = matrix_np.astype(np.complex128, copy=False)
        self.generator = self._numpy_generator
        self.hilbert_space = hilbert_space
        self.atol = float(atol)
        if use_scipy is None:
            has_matrix_exp = _has_backend_matrix_exp(self.backend)
            if has_matrix_exp:
                self._use_scipy = False
            elif _scipy_expm is not None:
                self._use_scipy = True
            else:
                raise RuntimeError(
                    "Backend lacks matrix_exp and SciPy is unavailable for fallback."
                )
        else:
            if use_scipy and _scipy_expm is None:
                raise RuntimeError("SciPy expm requested but SciPy is not available.")
            self._use_scipy = bool(use_scipy and _scipy_expm is not None)

        self._identity_backend = ensure_array(
            np.eye(hilbert_space.dimension, dtype=np.complex128),
            backend=self.backend,
        )
        self._last_contractivity_gap = float("nan")
        if ensure_contractive:
            eigenvalues_backend, _ = self.backend.eig(self._generator_backend)
            eigenvalues = ensure_numpy(eigenvalues_backend, backend=self.backend)
            if np.max(eigenvalues.real) > self.atol:
                raise TNFRValueError(
                    "ΔNFR generator is not contractive: positive real eigenvalues detected.",
                    context={"max_real_eigenvalue": np.max(eigenvalues.real)},
                    suggestion="Ensure generator is dissipative.",
                )

    def _propagator_backend(self, dt: float) -> Any:
        if self._use_scipy and _scipy_expm is not None:
            return ensure_array(
                _scipy_expm(dt * self._numpy_generator),
                backend=self.backend,
            )
        return self.backend.matrix_exp(dt * self._generator_backend)

    def frobenius_norm(
        self,
        density: Sequence[Sequence[complex]] | np.ndarray | Any,
        *,
        center: bool = False,
    ) -> float:
        """Return the Frobenius norm associated with the Hilbert space."""

        matrix = ensure_array(density, dtype=np.complex128, backend=self.backend)
        if matrix.shape != (self.hilbert_space.dimension, self.hilbert_space.dimension):
            raise TNFRValueError(
                "Density operator dimension mismatch.",
                context={
                    "expected_dimension": (
                        self.hilbert_space.dimension,
                        self.hilbert_space.dimension,
                    ),
                    "received_shape": matrix.shape,
                },
                suggestion="Ensure density operator matches Hilbert space dimension.",
            )
        if center:
            trace_value = _trace(matrix, backend=self.backend)
            trace_backend = trace_value.backend / self.hilbert_space.dimension
            matrix = matrix - trace_backend * self._identity_backend
        norm_backend = self.backend.norm(matrix, ord="fro")
        return float(np.asarray(ensure_numpy(norm_backend, backend=self.backend)))

    @property
    def last_contractivity_gap(self) -> float:
        """Return the latest monitored contractivity gap (NaN if unavailable)."""

        return float(self._last_contractivity_gap)

    def step(
        self,
        density: Sequence[Sequence[complex]] | np.ndarray | Any,
        *,
        dt: float = 1.0,
        normalize_trace: bool = True,
        enforce_contractivity: bool = True,
        raise_on_violation: bool = False,
        symmetrize: bool = True,
    ) -> Any:
        """Advance ``density`` by ``dt`` enforcing trace and contractivity control."""

        matrix = ensure_array(density, dtype=np.complex128, backend=self.backend)
        dim = self.hilbert_space.dimension
        if matrix.shape != (dim, dim):
            raise TNFRValueError(
                "Density operator dimension mismatch.",
                context={
                    "expected_dimension": (dim, dim),
                    "received_shape": matrix.shape,
                },
                suggestion="Ensure density operator matches Hilbert space dimension.",
            )

        initial_norm = None
        if enforce_contractivity:
            trace_value = _trace(matrix, backend=self.backend)
            trace_backend = trace_value.backend / dim
            centered = matrix - trace_backend * self._identity_backend
            initial_norm_backend = self.backend.norm(centered, ord="fro")
            try:
                initial_norm = float(
                    np.asarray(ensure_numpy(initial_norm_backend, backend=self.backend))
                )
            except (ValueError, TypeError, Exception):
                initial_norm = None

        vector = _vectorize_density(matrix, backend=self.backend)
        propagator = self._propagator_backend(dt)
        evolved_vec = self.backend.matmul(propagator, vector)
        evolved = _devectorize_density(evolved_vec, dim, backend=self.backend)

        if symmetrize:
            evolved = 0.5 * (evolved + self.backend.conjugate_transpose(evolved))

        if normalize_trace:
            trace_value = _trace(evolved, backend=self.backend)
            if trace_value.numpy is not None:
                if np.isclose(trace_value.numpy, 0.0, atol=self.atol):
                    raise TNFRValueError(
                        "Trace collapsed below tolerance during evolution.",
                        context={"trace": trace_value.numpy, "atol": self.atol},
                        suggestion="Check generator properties or initial state.",
                    )
                if not np.isclose(trace_value.numpy, 1.0, atol=10 * self.atol):
                    evolved = evolved / trace_value.backend
            else:
                # Tracing fallback: always normalize
                evolved = evolved / trace_value.backend

        if enforce_contractivity and initial_norm is not None:
            trace_value = _trace(evolved, backend=self.backend)
            trace_backend = trace_value.backend / dim
            centered = evolved - trace_backend * self._identity_backend
            evolved_norm_backend = self.backend.norm(centered, ord="fro")
            try:
                evolved_norm = float(
                    np.asarray(ensure_numpy(evolved_norm_backend, backend=self.backend))
                )
                self._last_contractivity_gap = initial_norm - evolved_norm
                if raise_on_violation and self._last_contractivity_gap < -5 * self.atol:
                    raise TNFRValueError(
                        "Contractivity violated: Frobenius norm increased beyond tolerance.",
                        context={
                            "initial_norm": initial_norm,
                            "evolved_norm": evolved_norm,
                            "gap": self._last_contractivity_gap,
                            "atol": self.atol,
                        },
                        suggestion="Ensure generator is contractive.",
                    )
            except (ValueError, TypeError, Exception):
                self._last_contractivity_gap = float("nan")
        else:
            self._last_contractivity_gap = float("nan")

        return evolved

    def evolve(
        self,
        density: Sequence[Sequence[complex]] | np.ndarray | Any,
        *,
        steps: int,
        dt: float = 1.0,
        normalize_trace: bool = True,
        enforce_contractivity: bool = True,
        raise_on_violation: bool = False,
        symmetrize: bool = True,
    ) -> Any:
        """Return trajectory of density operators for the contractive semigroup."""

        if steps < 0:
            raise TNFRValueError(
                "steps must be non-negative.",
                context={"steps": steps},
                suggestion="Provide a non-negative number of steps.",
            )

        current = ensure_array(density, dtype=np.complex128, backend=self.backend)
        dim = self.hilbert_space.dimension
        if current.shape != (dim, dim):
            raise TNFRValueError(
                "Density operator dimension mismatch.",
                context={"expected_shape": (dim, dim), "actual_shape": current.shape},
                suggestion="Ensure density operator matches Hilbert space dimension.",
            )

        trajectory: list[Any] = [current]
        for _ in range(steps):
            current = self.step(
                current,
                dt=dt,
                normalize_trace=normalize_trace,
                enforce_contractivity=enforce_contractivity,
                raise_on_violation=raise_on_violation,
                symmetrize=symmetrize,
            )
            trajectory.append(current)
        return self.backend.stack(trajectory, axis=0)