"""Successive Randomized Compression of a stack of tensor trains.

A stack ``T_1 . T_2 . ... . T_m`` is an ordered sequence of trains contracted
along their physical legs and compressed in a single SRC sweep. `apply` and
`compress` are the two-train and one-train special cases.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from ._sweep import sweep
from ._tensor_train import (
    MIN_SRC_SITES,
    check_exact_supported,
    exact_stack,
    normalize_stack,
)
from ._validation import sketch_size, validate_chi_out, validate_cutoff
from .utils import default_rng, get_xp, sketch_dtype

if TYPE_CHECKING:
    from collections.abc import Sequence

    import numpy as np
    from numpy.typing import DTypeLike

    from ._plan import Resources
    from ._tensor_train import Site

logger = logging.getLogger(__name__)

LOG_WARN_SMALL = (
    "The current SRC implementation targets tensor networks with 3 or more sites. "
    "Defaulting to an exact SVD-based contraction-compression."
)


def src(
    *trains: Sequence[Site],
    chi_out: int | np.integer,
    cutoff: float = 0.0,
    oversample: float | None = 1.0,
    dtype: DTypeLike | None = None,
    seed: int | None = None,
    device: str = "cpu",
    resources: Resources | None = None,
) -> list[np.ndarray]:
    """Contract a stack of tensor trains and compress the result with SRC.

    The stack ``src(T_1, T_2, ..., T_m)`` is the product ``T_1 T_2 ... T_m`` in
    mathematical order: ``T_m`` acts first. Trains are plain lists of per-site
    arrays in the default `quimb` layout, their kind inferred from the rank of the
    first site tensor. Each contraction joins the ``d`` leg of a train with the
    ``u`` (or physical) leg of the next. Accepted stacks:

    - one MPS or one MPO: plain compression;
    - ``A_1, ..., A_k``: an MPO product, giving an MPO;
    - ``A_1, ..., A_k, psi``: an MPO product applied to a ket, giving an MPS;
    - ``phi, A_1, ..., A_k``: a bra times an MPO product, giving an MPS on the
      ``d`` leg of ``A_k``.

    A leading MPS is contracted without conjugation, as the row vector
    ``phi^T A_1 ... A_k``. For the physical bra ``