"""ΔNFR generator construction utilities."""

from __future__ import annotations

from typing import Final, Sequence

from numpy.random import Generator

from .._exact_time import finite_represented_real
from .._spectral_expectation import positive_spectral_dimension
from ..config.parsing import parse_bool
from ._complex_arrays import (
    backend_complex_array,
    nonnegative_tolerance,
    numpy_complex_array,
)
from .backend import ensure_numpy, get_backend
from .unified_numerical import TNFRValueError, np

__all__ = ["build_delta_nfr", "build_lindblad_delta_nfr"]

_TOPOLOGIES: Final[set[str]] = {"laplacian", "adjacency"}


def _generator_scale(nu_f: float, scale: float) -> float:
    """Admit signed real factors and their represented product."""
    try:
        frequency = finite_represented_real(nu_f, "nu_f")[0]
        amplitude = finite_represented_real(scale, "scale")[0]
        product = finite_represented_real(frequency * amplitude, "nu_f * scale")[0]
    except (TypeError, ValueError, OverflowError) as exc:
        raise TNFRValueError(
            "Generator scaling must be finite representable real"
        ) from exc
    if frequency != 0 and amplitude != 0 and product == 0:
        raise TNFRValueError("Nonzero generator scaling product underflows to zero")
    return product


def _ring_adjacency(dim: int) -> np.ndarray:
    """Return the adjacency matrix for a coherent ring topology."""

    adjacency: np.ndarray = np.zeros((dim, dim), dtype=float)
    if dim == 1:
        return adjacency

    indices = np.arange(dim)
    adjacency[indices, (indices + 1) % dim] = 1.0
    adjacency[(indices + 1) % dim, indices] = 1.0
    return adjacency


def _laplacian_from_adjacency(adjacency: np.ndarray) -> np.ndarray:
    """Construct a Laplacian operator from an adjacency matrix."""

    degrees = adjacency.sum(axis=1)
    laplacian = np.diag(degrees) - adjacency
    return laplacian


def _hermitian_noise(dim: int, rng: Generator) -> np.ndarray:
    """Generate a Hermitian noise matrix with reproducible statistics."""

    real = rng.standard_normal((dim, dim))
    imag = rng.standard_normal((dim, dim))
    noise = real + 1j * imag
    return 0.5 * (noise + noise.conj().T)


def _as_square_matrix(
    matrix: Sequence[Sequence[complex]] | np.ndarray,
    *,
    expected_dim: int | None = None,
    label: str = "matrix",
) -> np.ndarray:
    """Return ``matrix`` as a square :class:`numpy.ndarray` with validation."""

    array = numpy_complex_array(matrix, label=label)
    if array.ndim != 2 or array.shape[0] != array.shape[1]:
        raise TNFRValueError(
            f"{label} must be a square matrix.",
            context={"shape": array.shape},
            suggestion="Provide a square matrix.",
        )
    if expected_dim is not None and array.shape[0] != expected_dim:
        raise TNFRValueError(
            f"{label} dimension mismatch.",
            context={"expected": expected_dim, "received": array.shape[0]},
            suggestion="Ensure matrix dimension matches expected dimension.",
        )
    return array


def build_delta_nfr(
    dim: int,
    *,
    topology: str = "laplacian",
    nu_f: float = 1.0,
    scale: float = 1.0,
    rng: Generator | None = None,
) -> np.ndarray:
    """Construct a Hermitian ΔNFR generator using canonical TNFR topologies.

    Parameters
    ----------
    dim:
        Dimensionality of the Hilbert space supporting the ΔNFR operator.
    topology:
        Requested canonical topology. Supported values are ``"laplacian"``
        and ``"adjacency"``.
    nu_f:
        Finite represented real auxiliary frequency multiplier. Signed values
        here do not define an admitted nodal capacity.
    scale:
        Additional finite represented real amplitude factor. Signed factors
        are admitted; their product must also be representable.
    rng:
        Optional NumPy :class:`~numpy.random.Generator` used to inject
        reproducible Hermitian noise.
    """

    dim = positive_spectral_dimension(dim, label="Generator dimension")
    scaling = _generator_scale(nu_f, scale)

    if topology not in _TOPOLOGIES:
        allowed = ", ".join(sorted(_TOPOLOGIES))
        raise TNFRValueError(
            f"Unknown ΔNFR topology: {topology}.",
            context={"topology": topology, "allowed": list(_TOPOLOGIES)},
            suggestion=f"Expected one of: {allowed}.",
        )

    adjacency = _ring_adjacency(dim)
    if topology == "laplacian":
        base = _laplacian_from_adjacency(adjacency)
    else:
        base = adjacency

    matrix: np.ndarray = base.astype(np.complex128, copy=False)

    if rng is not None:
        noise = _hermitian_noise(dim, rng)
        matrix = matrix + (1.0 / np.sqrt(dim)) * noise

    with np.errstate(over="ignore", invalid="ignore"):
        matrix *= scaling
    # Every summand is Hermitian and the admitted scale is real; averaging
    # equal adjoint entries again would only risk avoidable overflow.
    backend = get_backend()
    return np.asarray(
        ensure_numpy(
            backend_complex_array(matrix, backend=backend, label="Generated operator"),
            backend=backend,
        ),
        dtype=np.complex128,
    )


def build_lindblad_delta_nfr(
    *,
    hamiltonian: Sequence[Sequence[complex]] | np.ndarray | None = None,
    collapse_operators: (
        Sequence[Sequence[Sequence[complex]] | np.ndarray] | None
    ) = None,
    dim: int | None = None,
    nu_f: float = 1.0,
    scale: float = 1.0,
    ensure_trace_preserving: bool = True,
    ensure_contractive: bool = True,
    atol: float = 1e-9,
) -> np.ndarray:
    """Construct a Lindblad ΔNFR generator in Liouville space.

    The resulting matrix acts on vectorised density operators using the
    canonical column-major flattening.  The construction follows the standard
    Gorini–Kossakowski–Sudarshan–Lindblad prescription while exposing TNFR
    semantics through ``ν_f`` and ``scale``.

    Parameters
    ----------
    hamiltonian:
        Optional coherent component.  When ``None`` a null Hamiltonian is
        assumed. Matrix entries must be finite real or complex scalars with
        representable components; Boolean and textual entries are rejected.
    collapse_operators:
        Iterable with the dissipative operators driving the contractive
        semigroup.  Each entry must be square with the same dimension as the
        Hamiltonian and obey the same scalar admission. When ``None`` the
        generator reduces to the coherent part.
    dim:
        Explicit Hilbert-space dimension.  Only required if neither
        ``hamiltonian`` nor ``collapse_operators`` are provided.  When supplied,
        it must match the dimension inferred from the Hamiltonian and collapse
        operators.
    nu_f, scale:
        Finite represented real factors applied uniformly to the generator.
        Their product must be representable. Signed factors are admitted;
        a forward GKSL semigroup requires a nonnegative product.
        These auxiliary multipliers do not admit negative nodal capacities.
    ensure_trace_preserving:
        When ``True`` (default) the resulting superoperator is validated to
        annihilate the left identity row, preserving trace within tolerance.
        This does not require unitality: the identity may evolve.
    ensure_contractive:
        When ``True`` (default) the spectrum is required to have non-positive
        real parts within ``atol``. This spectral check alone does not certify
        complete positivity or contraction in every norm.
    atol:
        Finite nonnegative represented tolerance for Hermiticity, trace and
        spectral checks.
    """

    scaling = _generator_scale(nu_f, scale)
    atol = nonnegative_tolerance(atol)
    ensure_trace_preserving = parse_bool(ensure_trace_preserving)
    ensure_contractive = parse_bool(ensure_contractive)
    if dim is not None:
        dim = positive_spectral_dimension(dim, label="Generator dimension")
    operators = [] if collapse_operators is None else list(collapse_operators)

    hermitian = None
    dissipators = []
    inferred_dim: int | None = dim
    if hamiltonian is not None:
        hermitian = _as_square_matrix(hamiltonian, label="hamiltonian")
        inferred_dim = hermitian.shape[0]
    elif operators:
        dissipators.append(
            _as_square_matrix(operators[0], label="collapse operator[0]")
        )
        inferred_dim = dissipators[0].shape[0]

    if inferred_dim is None:
        raise TNFRValueError(
            "dim must be supplied when no operators are provided.",
            context={"dim": dim, "hamiltonian": hamiltonian, "operators": operators},
            suggestion="Provide explicit dimension or operators.",
        )

    if inferred_dim <= 0:
        raise TNFRValueError(
            "ΔNFR generators require a positive dimension.",
            context={"inferred_dim": inferred_dim},
            suggestion="Ensure dimension is positive.",
        )

    dimension = inferred_dim

    if dim is not None and dim != dimension:
        raise TNFRValueError(
            "Provided dim is inconsistent with the supplied operators.",
            context={"expected": dimension, "received": dim},
            suggestion="Ensure explicit dimension matches operators.",
        )

    if hermitian is None:
        hermitian = np.zeros((dimension, dimension), dtype=np.complex128)
    else:
        if not np.allclose(hermitian, hermitian.conj().T, atol=atol):
            raise TNFRValueError(
                "Hamiltonian component must be Hermitian within tolerance.",
                context={"atol": atol},
                suggestion="Ensure Hamiltonian is Hermitian.",
            )

    first_unread = len(dissipators)
    dissipators.extend(
        _as_square_matrix(
            operator, expected_dim=dimension, label=f"collapse operator[{index}]"
        )
        for index, operator in enumerate(operators[first_unread:], start=first_unread)
    )

    identity = np.eye(dimension, dtype=np.complex128)
    liouvillian = -1j * (np.kron(identity, hermitian) - np.kron(hermitian.T, identity))

    for operator in dissipators:
        adjoint_product = operator.conj().T @ operator
        liouvillian += np.kron(operator.conj(), operator)
        liouvillian -= 0.5 * np.kron(identity, adjoint_product)
        liouvillian -= 0.5 * np.kron(adjoint_product.T, identity)

    with np.errstate(over="ignore", invalid="ignore"):
        liouvillian *= scaling
    liouvillian = numpy_complex_array(liouvillian, label="Generated Lindblad operator")

    if ensure_trace_preserving:
        identity_vec = identity.reshape(dimension * dimension, order="F")
        left_residual = identity_vec.conj().T @ liouvillian
        if not np.allclose(left_residual, np.zeros_like(left_residual), atol=10 * atol):
            raise TNFRValueError(
                "Lindblad generator must preserve the trace of density operators.",
                context={"residual_norm": np.linalg.norm(left_residual)},
                suggestion="Check dissipator construction.",
            )

    backend = get_backend()
    liouvillian_backend = backend_complex_array(
        liouvillian, backend=backend, label="Generated Lindblad operator"
    )

    if ensure_contractive:
        eigenvalues_backend, _ = backend.eig(liouvillian_backend)
        eigenvalues = numpy_complex_array(
            ensure_numpy(eigenvalues_backend, backend=backend),
            label="Generator spectrum",
        )
        if np.max(eigenvalues.real) > atol:
            raise TNFRValueError(
                "Lindblad generator is not contractive: spectrum has positive real components.",
                context={"max_real_eigenvalue": np.max(eigenvalues.real)},
                suggestion="Ensure generator is dissipative.",
            )

    return np.asarray(
        ensure_numpy(liouvillian_backend, backend=backend), dtype=np.complex128
    )