import math
import torch
from typing import List, Optional, Tuple


class TensorTrain:
    """
    Minimal tensor train representation with d degrees of freedom.

    Each core is a rank-3 tensor with shape (r_{k-1}, n_k, r_k) where
    r_0 = 1 and r_d = 1 for consistency.
    """

    def __init__(
        self,
        dims: List[int],
        ranks: Optional[List[int]] = None,
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ):
        """
        Initialize a tensor train.

        Args:
            dims: List of dimensions [n_1, n_2, ..., n_d]
            ranks: List of internal ranks [r_1, r_2, ..., r_{d-1}].
                   If None, defaults to all ones (rank-1 TT).
            dtype: Torch dtype for the tensors
            device: Torch device
        """
        self.d = len(dims)
        self.dims = list(dims)
        self.dtype = dtype
        self.device = device

        # Build full rank list: [r_0=1, r_1, ..., r_{d-1}, r_d=1]
        if ranks is None:
            ranks = [1] * (self.d - 1)

        if len(ranks) != self.d - 1:
            raise ValueError(
                f"Expected {self.d - 1} internal ranks, got {len(ranks)}"
            )

        self.ranks = [1] + list(ranks) + [1]

        # Initialize cores: core[k] has shape (r_{k-1}, n_k, r_k)
        self.cores: List[torch.Tensor] = []
        for k in range(self.d):
            shape = (self.ranks[k], self.dims[k], self.ranks[k + 1])
            core = torch.zeros(shape, dtype=dtype, device=device)
            self.cores.append(core)

    @classmethod
    def random(
        cls,
        dims: List[int],
        ranks: List[int],
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ) -> "TensorTrain":
        """Create a tensor train with random Gaussian entries."""
        tt = cls(dims, ranks, dtype, device)
        for k in range(tt.d):
            tt.cores[k] = torch.randn_like(tt.cores[k])
        return tt

    @classmethod
    def ones(
        cls,
        dims: List[int],
        value: float = 1.0,
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ) -> "TensorTrain":
        """
        Create a rank-1 tensor train representing a constant tensor.

        Every element of the full tensor equals `value`.

        Args:
            dims: Dimensions [n_1, n_2, ..., n_d]
            value: Constant value for all elements (default 1.0)
            dtype: Torch dtype
            device: Torch device

        Returns:
            Rank-1 TensorTrain where to_tensor() gives all `value`
        """
        tt = cls(dims, None, dtype, device)  # rank-1 by default
        for k in range(tt.d):
            tt.cores[k] = torch.ones_like(tt.cores[k])
        # Scale the first core by value
        tt.cores[0] = value * tt.cores[0]
        return tt

    @classmethod
    def zeros(
        cls,
        dims: List[int],
        ranks: Optional[List[int]] = None,
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ) -> "TensorTrain":
        """
        Create a tensor train with all zero entries.

        Args:
            dims: Dimensions [n_1, n_2, ..., n_d]
            ranks: Internal ranks (default: rank-1)
            dtype: Torch dtype
            device: Torch device

        Returns:
            TensorTrain where to_tensor() gives all zeros
        """
        # The default __init__ already creates zero-initialized cores
        return cls(dims, ranks, dtype, device)

    @classmethod
    def from_dense(
        cls,
        tensor: torch.Tensor,
        max_rank: Optional[int] = None,
        rel_tol: Optional[float] = None,
        abs_tol: Optional[float] = None,
    ) -> "TensorTrain":
        """
        Convert a dense tensor to TensorTrain format using TT-SVD.

        Performs sequential SVD decomposition from left to right,
        optionally truncating ranks based on specified tolerances.

        Args:
            tensor: Dense tensor of any shape (n_1, n_2, ..., n_d)
            max_rank: Maximum rank at each bond (None = no limit)
            rel_tol: Relative tolerance - keep singular values where
                     s[i] >= rel_tol * s[0]
            abs_tol: Absolute tolerance - keep singular values where
                     s[i] >= abs_tol

        Returns:
            TensorTrain representation of the input tensor

        Example:
            dense = torch.randn(4, 5, 6)
            tt = TensorTrain.from_dense(dense)  # Exact representation
            tt_approx = TensorTrain.from_dense(dense, max_rank=3)  # Truncated
        """
        if tensor.ndim == 0:
            raise ValueError("Cannot convert scalar to TensorTrain")

        dims = list(tensor.shape)
        d = len(dims)
        dtype = tensor.dtype
        device = tensor.device

        if d == 1:
            # Special case: 1D tensor -> single core (1, n, 1)
            tt = cls(dims, [], dtype, device)
            tt.cores[0] = tensor.reshape(1, dims[0], 1)
            return tt

        cores = []
        current = tensor

        # Left-to-right sweep: extract cores one at a time
        for k in range(d - 1):
            # current shape: (r_left, n_k, n_{k+1}, ..., n_d) for k > 0
            # or (n_0, n_1, ..., n_d) for k = 0
            if k == 0:
                r_left = 1
                n_k = dims[k]
                remaining = current.reshape(n_k, -1)  # (n_k, prod(rest))
            else:
                r_left = current.shape[0]
                n_k = dims[k]
                # current: (r_left, n_k, remaining_prod)
                remaining = current.reshape(r_left * n_k, -1)

            # SVD: (r_left * n_k, remaining) -> U @ S @ Vt
            U, S, Vt = torch.linalg.svd(remaining, full_matrices=False)

            # Determine rank to keep
            r_new = len(S)

            if max_rank is not None:
                r_new = min(r_new, max_rank)

            if rel_tol is not None and len(S) > 0:
                threshold = rel_tol * S[0]
                r_new = min(r_new, int((S >= threshold).sum().item()))

            if abs_tol is not None:
                r_new = min(r_new, int((S >= abs_tol).sum().item()))

            # Ensure at least rank 1
            r_new = max(1, r_new)

            # Truncate
            U = U[:, :r_new]
            S = S[:r_new]
            Vt = Vt[:r_new, :]

            # Core k: (r_left, n_k, r_new)
            if k == 0:
                core = U.reshape(1, n_k, r_new)
            else:
                core = U.reshape(r_left, n_k, r_new)
            cores.append(core)

            # Propagate S @ Vt to next iteration
            current = (S.unsqueeze(1) * Vt)  # (r_new, remaining_prod)

        # Last core: (r_left, n_d, 1)
        r_left = current.shape[0]
        n_d = dims[-1]
        cores.append(current.reshape(r_left, n_d, 1))

        # Build TensorTrain
        internal_ranks = [c.shape[2] for c in cores[:-1]]
        tt = cls(dims, internal_ranks, dtype, device)
        for k in range(d):
            tt.cores[k] = cores[k]

        return tt

    @property
    def shape(self) -> Tuple[int, ...]:
        """Logical shape of the tensor: (n_1, n_2, ..., n_d)."""
        return tuple(self.dims)

    @property
    def ttshape(self) -> List[Tuple[int, int, int]]:
        """Shape of each TT core: [(r_0, n_1, r_1), (r_1, n_2, r_2), ...]."""
        return [
            (self.ranks[k], self.dims[k], self.ranks[k + 1])
            for k in range(self.d)
        ]

    @property
    def numel(self) -> int:
        """Logical number of elements: n_1 * n_2 * ... * n_d."""
        result = 1
        for n in self.dims:
            result *= n
        return result

    @property
    def ttnumel(self) -> int:
        """Actual storage size: sum of r_{k-1} * n_k * r_k over all cores."""
        return sum(
            self.ranks[k] * self.dims[k] * self.ranks[k + 1]
            for k in range(self.d)
        )

    @property
    def ndim(self) -> int:
        """Number of dimensions (same as d)."""
        return self.d

    def __getitem__(self, key):
        """
        Get core, element, or slice from the tensor train.

        Three modes:
        1. Core access: tt[k] returns the k-th core tensor
        2. Element access: tt[i, j, k] returns a scalar element
        3. Slicing: tt[:, 1:3, 4] returns a new TT with reduced dimensions

        Args:
            key: int for core access, tuple for element/slice access

        Returns:
            Core tensor, scalar, or new TensorTrain depending on key type

        Examples:
            tt[0]           # First core tensor
            tt[1, 2, 3]     # Element at index (1, 2, 3)
            tt[:, 2, :]     # Slice selecting index 2 from middle dimension
            tt[0:2, :, 1:]  # Slice with ranges
        """
        # Single int: core access (backward compatible)
        if isinstance(key, int):
            return self.cores[key]

        # Convert to tuple if needed
        if not isinstance(key, tuple):
            key = (key,)

        # Expand ellipsis and pad with slice(None)
        key = self._expand_key(key)

        # Check if all integers (pure element access)
        if all(isinstance(k, int) for k in key):
            return self._get_element(key)

        # Mixed: slicing
        return self._get_slice(key)

    def _expand_key(self, key: tuple) -> tuple:
        """Expand ellipsis and pad key to length d."""
        # Handle Ellipsis
        if Ellipsis in key:
            idx = key.index(Ellipsis)
            before = key[:idx]
            after = key[idx + 1:]
            n_fill = self.d - len(before) - len(after)
            key = before + (slice(None),) * n_fill + after

        # Pad with slice(None) if too short
        if len(key) < self.d:
            key = key + (slice(None),) * (self.d - len(key))

        if len(key) != self.d:
            raise IndexError(
                f"Too many indices: got {len(key)}, expected {self.d}"
            )

        return key

    def _get_element(self, key: tuple) -> torch.Tensor:
        """Get single element tt[i, j, k, ...] as a 0-d tensor."""
        # Start with first core, selecting index
        # core[0]: (1, n_0, r_1) -> select key[0] -> (1, r_1) -> squeeze to (r_1,)
        result = self.cores[0][0, key[0], :]

        # Contract through middle cores
        for k in range(1, self.d - 1):
            # core[k]: (r_k, n_k, r_{k+1}) -> select key[k] -> (r_k, r_{k+1})
            result = result @ self.cores[k][:, key[k], :]

        # Final core
        # core[d-1]: (r_{d-1}, n_{d-1}, 1) -> select key[d-1] -> (r_{d-1}, 1)
        result = result @ self.cores[-1][:, key[-1], :]

        # Return 0-d tensor (consistent with torch, preserves autograd)
        return result.squeeze()

    def _get_slice(self, key: tuple) -> "TensorTrain":
        """
        Get slice tt[:, j, 1:3, ...] returning a new TensorTrain.

        Integer indices contract that dimension (absorbing into neighbors).
        Slice indices keep that dimension (possibly reduced).
        """
        # Separate into contracted (int) and kept (slice) indices
        # We'll process left-to-right, accumulating contractions

        # First pass: slice each core and identify which are contracted
        sliced_cores = []
        is_contracted = []

        for k, idx in enumerate(key):
            core = self.cores[k]  # (r_l, n, r_r)

            if isinstance(idx, int):
                # Contract this dimension: (r_l, n, r_r)[:, idx, :] -> (r_l, r_r)
                sliced_cores.append(core[:, idx, :])
                is_contracted.append(True)
            elif isinstance(idx, slice):
                # Keep this dimension (possibly sliced)
                sliced_cores.append(core[:, idx, :])
                is_contracted.append(False)
            else:
                raise TypeError(f"Invalid index type: {type(idx)}")

        # If all contracted, return scalar (shouldn't happen, _get_element handles this)
        if all(is_contracted):
            return self._get_element(key)

        # Absorb contracted cores into kept cores
        # Strategy: left-to-right pass absorbs from left, right-to-left absorbs from right

        # Left-to-right: accumulate left contractions
        left_matrix = None
        result_cores = []

        for k in range(self.d):
            core = sliced_cores[k]

            if is_contracted[k]:
                # This is a matrix (r_l, r_r), accumulate it
                if left_matrix is None:
                    left_matrix = core
                else:
                    left_matrix = left_matrix @ core
            else:
                # This is a kept core (r_l, n', r_r)
                if left_matrix is not None:
                    # Absorb accumulated left matrix
                    # left_matrix: (r_start, r_l), core: (r_l, n', r_r)
                    # result: (r_start, n', r_r)
                    r_l, n_prime, r_r = core.shape
                    core = torch.einsum('ij,jkl->ikl', left_matrix, core)
                    left_matrix = None

                result_cores.append(core)

        # If there's remaining left_matrix, absorb into last kept core
        if left_matrix is not None and result_cores:
            last_core = result_cores[-1]
            # last_core: (r_l, n', r_r), left_matrix: (r_r, r_end)
            # result: (r_l, n', r_end)
            result_cores[-1] = torch.einsum('ijk,kl->ijl', last_core, left_matrix)

        # Build new TensorTrain
        new_dims = [core.shape[1] for core in result_cores]
        new_ranks = [core.shape[2] for core in result_cores[:-1]]

        result = TensorTrain(new_dims, new_ranks if new_ranks else None,
                            self.dtype, self.device)
        for k, core in enumerate(result_cores):
            result.cores[k] = core

        return result

    def __setitem__(self, k: int, value: torch.Tensor):
        """Set the k-th core tensor."""
        expected_shape = (self.ranks[k], self.dims[k], self.ranks[k + 1])
        if value.shape != expected_shape:
            raise ValueError(
                f"Core {k} expected shape {expected_shape}, got {value.shape}"
            )
        self.cores[k] = value

    def __len__(self) -> int:
        """Return the number of cores (degrees of freedom)."""
        return self.d

    def __add__(self, other: "TensorTrain") -> "TensorTrain":
        """Add two TensorTrains: self + other."""
        if isinstance(other, TensorTrain):
            return add(self, other)
        return NotImplemented

    def __sub__(self, other: "TensorTrain") -> "TensorTrain":
        """Subtract two TensorTrains: self - other."""
        if isinstance(other, TensorTrain):
            return add(self, scale(other, -1.0))
        return NotImplemented

    def __mul__(self, scalar: float) -> "TensorTrain":
        """Scalar multiplication: self * scalar."""
        if isinstance(scalar, (int, float)):
            return scale(self, float(scalar))
        return NotImplemented

    def __rmul__(self, scalar: float) -> "TensorTrain":
        """Scalar multiplication: scalar * self."""
        if isinstance(scalar, (int, float)):
            return scale(self, float(scalar))
        return NotImplemented

    def __neg__(self) -> "TensorTrain":
        """Negation: -self."""
        return scale(self, -1.0)

    def total_params(self) -> int:
        """Return total number of parameters in the TT (alias for ttnumel)."""
        return self.ttnumel

    def copy(self) -> "TensorTrain":
        """Return a deep copy of this tensor train."""
        tt = TensorTrain(self.dims, self.ranks[1:-1], self.dtype, self.device)
        for k in range(self.d):
            tt.cores[k] = self.cores[k].clone()
        return tt

    def clone(self) -> "TensorTrain":
        """Return a deep copy of this tensor train (torch-like API)."""
        return self.copy()

    def detach(self) -> "TensorTrain":
        """
        Detach tensor train from computation graph (torch-like API).

        Returns a new TensorTrain with all cores detached from the
        autograd computation graph. The new cores will not require
        gradients.

        Returns:
            New TensorTrain with detached cores
        """
        tt = TensorTrain(self.dims, self.ranks[1:-1], self.dtype, self.device)
        for k in range(self.d):
            tt.cores[k] = self.cores[k].detach()
        return tt

    def to(
        self,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
    ) -> "TensorTrain":
        """
        Convert tensor train to different dtype and/or device (torch-like API).

        Args:
            dtype: Target dtype (if None, keeps current dtype)
            device: Target device (if None, keeps current device)

        Returns:
            New TensorTrain with converted cores
        """
        if dtype is None:
            dtype = self.dtype
        if device is None:
            device = self.device

        tt = TensorTrain(self.dims, self.ranks[1:-1], dtype, device)
        for k in range(self.d):
            tt.cores[k] = self.cores[k].to(dtype=dtype, device=device)
        return tt

    def tt_shape(self) -> List[Tuple[int, int, int]]:
        """
        Return the shapes of all TT cores.

        Returns:
            List of tuples [(r0, n1, r1), (r1, n2, r2), ..., (r_{d-1}, nd, rd)]
        """
        return [tuple(core.shape) for core in self.cores]

    def tt_ranks(self) -> Tuple[int, ...]:
        """
        Return the TT ranks including boundary ranks.

        Returns:
            Tuple of ranks (r0, r1, ..., rd) where r0 = rd = 1
        """
        return tuple(self.ranks)

    def ncores(self) -> int:
        """Return the number of cores in the TT."""
        return self.d

    def tt_cores(self) -> List[torch.Tensor]:
        """
        Return the list of TT cores.

        Returns:
            List of core tensors [core0, core1, ..., core_{d-1}]
        """
        return self.cores

    def storage_numel(self) -> int:
        """
        Return the total number of stored elements across all cores.

        Returns:
            Sum of numel() for all core tensors
        """
        return self.total_params()

    def compression_ratio(self) -> float:
        """
        Return the compression ratio: logical_numel / storage_numel.

        Higher values indicate better compression.

        Returns:
            Ratio of logical tensor size to stored tensor size
        """
        return self.numel / self.storage_numel()

    def merge_modes(self, start: int, end: int) -> "TensorTrain":
        """
        Merge adjacent modes [start, end) into a single mode.

        Args:
            start: First mode to merge (inclusive)
            end: Last mode to merge (exclusive)

        Returns:
            New TensorTrain with merged modes

        Example:
            tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
            tt2 = tt.merge_modes(0, 2)  # dims becomes [6, 4]
        """
        if start < 0 or end > self.d or start >= end:
            raise ValueError(
                f"Invalid mode range [{start}, {end}) for TT with {self.d} modes"
            )

        if end - start == 1:
            return self.clone()

        # Contract cores from start to end-1
        merged_core = self.cores[start]  # (r_left, n_start, r_1)
        for k in range(start + 1, end):
            # merged_core: (r_left, n_merged, r_k)
            # cores[k]: (r_k, n_k, r_{k+1})
            r_left = merged_core.shape[0]
            n_merged = merged_core.shape[1]
            r_mid = merged_core.shape[2]
            n_k = self.cores[k].shape[1]
            r_right = self.cores[k].shape[2]

            # Contract: (r_left, n_merged, r_mid) x (r_mid, n_k, r_right)
            # -> (r_left, n_merged, n_k, r_right) -> (r_left, n_merged*n_k, r_right)
            merged_core = torch.einsum('imr,rjk->imjk', merged_core, self.cores[k])
            merged_core = merged_core.reshape(r_left, n_merged * n_k, r_right)

        # Build new dims and cores
        merged_dim = 1
        for k in range(start, end):
            merged_dim *= self.dims[k]

        new_dims = self.dims[:start] + [merged_dim] + self.dims[end:]
        new_ranks = self.ranks[1:start+1] + self.ranks[end:-1]

        result = TensorTrain(new_dims, new_ranks if new_ranks else None,
                             self.dtype, self.device)

        # Copy cores before merged region
        for k in range(start):
            result.cores[k] = self.cores[k].clone()

        # Set merged core
        result.cores[start] = merged_core

        # Copy cores after merged region
        for k in range(end, self.d):
            result.cores[start + 1 + (k - end)] = self.cores[k].clone()

        return result

    def flatten(self, start_dim: int = 0, end_dim: int = -1) -> "TensorTrain":
        """
        Flatten a contiguous range of dimensions into one.

        PyTorch-compatible API: both start_dim and end_dim are inclusive,
        and negative indexing is supported.

        Args:
            start_dim: First dimension to flatten (default: 0)
            end_dim: Last dimension to flatten, inclusive (default: -1, i.e., last)

        Returns:
            New TensorTrain with flattened dimensions

        Example:
            tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
            tt.flatten()        # shape becomes (120,) - flatten all
            tt.flatten(0, 1)    # shape becomes (6, 4, 5)
            tt.flatten(1, 2)    # shape becomes (2, 12, 5)
            tt.flatten(1, -1)   # shape becomes (2, 60)
        """
        # Handle negative indexing
        if start_dim < 0:
            start_dim = self.d + start_dim
        if end_dim < 0:
            end_dim = self.d + end_dim

        if start_dim < 0 or start_dim >= self.d:
            raise ValueError(
                f"Invalid start_dim={start_dim} for TT with {self.d} modes"
            )
        if end_dim < 0 or end_dim >= self.d:
            raise ValueError(
                f"Invalid end_dim={end_dim} for TT with {self.d} modes"
            )
        if start_dim > end_dim:
            raise ValueError(
                f"start_dim ({start_dim}) must be <= end_dim ({end_dim})"
            )

        # Convert to merge_modes convention (exclusive end)
        return self.merge_modes(start_dim, end_dim + 1)

    def split_mode(self, mode: int, shape: List[int]) -> "TensorTrain":
        """
        Split a single mode into multiple modes.

        The mode dimension must equal the product of shape dimensions.
        New cores are connected with rank-1 bonds.

        Args:
            mode: Index of mode to split
            shape: List of dimensions to split into

        Returns:
            New TensorTrain with split mode

        Example:
            tt = TensorTrain.random([6, 4], ranks=[5])
            tt2 = tt.split_mode(0, [2, 3])  # dims becomes [2, 3, 4]
        """
        if mode < 0 or mode >= self.d:
            raise ValueError(f"Invalid mode {mode} for TT with {self.d} modes")

        if math.prod(shape) != self.dims[mode]:
            raise ValueError(
                f"Product of shape {shape} = {math.prod(shape)} "
                f"doesn't match mode dimension {self.dims[mode]}"
            )

        if len(shape) == 1:
            return self.clone()

        # Get the core to split: (r_left, n, r_right)
        core = self.cores[mode]
        r_left, n, r_right = core.shape

        # Use SVD-based splitting to peel off dimensions one at a time
        split_cores = []
        current = core  # (r_left, n, r_right)
        current_r_left = r_left

        for i, f in enumerate(shape[:-1]):
            remaining_prod = math.prod(shape[i+1:])
            # current: (current_r_left, f * remaining_prod, r_right)
            # Reshape to (current_r_left * f, remaining_prod * r_right)
            mat = current.reshape(current_r_left * f, remaining_prod * r_right)

            # SVD to find optimal rank (or use full rank for exactness)
            U, S, Vh = torch.linalg.svd(mat, full_matrices=False)

            # Keep all singular values for exact representation
            rank = S.shape[0]
            U = U @ torch.diag(S)  # (current_r_left * f, rank)
            # Vh: (rank, remaining_prod * r_right)

            # Reshape U to core: (current_r_left, f, rank)
            new_core = U.reshape(current_r_left, f, rank)
            split_cores.append(new_core)

            # Update for next iteration
            current = Vh.reshape(rank, remaining_prod, r_right)
            current_r_left = rank

        # Last core: (current_r_left, shape[-1], r_right)
        split_cores.append(current)

        # Build new TensorTrain
        new_dims = self.dims[:mode] + list(shape) + self.dims[mode+1:]
        new_ranks = (
            self.ranks[1:mode+1] +
            [c.shape[2] for c in split_cores[:-1]] +
            self.ranks[mode+1:-1]
        )

        result = TensorTrain(new_dims, new_ranks if new_ranks else None,
                             self.dtype, self.device)

        # Copy cores before split
        for k in range(mode):
            result.cores[k] = self.cores[k].clone()

        # Set split cores
        for i, c in enumerate(split_cores):
            result.cores[mode + i] = c

        # Copy cores after split
        for k in range(mode + 1, self.d):
            result.cores[mode + len(shape) - 1 + (k - mode)] = self.cores[k].clone()

        return result

    def reshape(self, new_shape: List[int]) -> "TensorTrain":
        """
        Reshape the TensorTrain to a new shape.

        Only supports reshapes that can be achieved by merging adjacent modes
        and/or splitting modes. Raises an error if the reshape would require
        permutation.

        Args:
            new_shape: Target shape (must have same total elements)

        Returns:
            New TensorTrain with target shape

        Example:
            tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
            tt.reshape([6, 4])    # merge first two modes
            tt.reshape([2, 12])   # merge last two modes
            tt.reshape([24])      # merge all modes
        """
        new_shape = list(new_shape)

        if self.numel != math.prod(new_shape):
            raise ValueError(
                f"Cannot reshape: {self.shape} ({self.numel} elements) -> "
                f"{tuple(new_shape)} ({math.prod(new_shape)} elements)"
            )

        if list(self.dims) == new_shape:
            return self.clone()

        # Try to find a sequence of merges and splits
        # Strategy: walk through old and new dims, matching products

        old_dims = list(self.dims)
        result = self.clone()

        # First pass: figure out the operations needed
        ops = _compute_reshape_ops(old_dims, new_shape)

        # Apply operations
        for op, args in ops:
            if op == 'merge':
                start, end = args
                result = result.merge_modes(start, end)
            elif op == 'split':
                mode, shape = args
                result = result.split_mode(mode, shape)

        return result

    def swap_adjacent(
        self,
        i: int,
        max_rank: Optional[int] = None,
        rel_tol: Optional[float] = None,
        abs_tol: Optional[float] = None,
    ) -> "TensorTrain":
        """
        Swap two adjacent modes i and i+1.

        Contracts cores[i] and cores[i+1], swaps the physical dimensions,
        and re-splits via SVD.

        Args:
            i: Index of first mode to swap (swaps modes i and i+1)
            max_rank: Maximum rank after SVD (None = no limit)
            rel_tol: Relative tolerance for SVD truncation
            abs_tol: Absolute tolerance for SVD truncation

        Returns:
            New TensorTrain with modes i and i+1 swapped

        Example:
            tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
            tt2 = tt.swap_adjacent(0)  # shape becomes (3, 2, 4)
            tt3 = tt.swap_adjacent(1)  # shape becomes (2, 4, 3)
        """
        if i < 0 or i >= self.d - 1:
            raise ValueError(
                f"Invalid index {i} for swap_adjacent: need 0 <= i < d-1 = {self.d - 1}"
            )

        # Get the two cores to swap
        # core_i: (r_left, n_i, r_mid)
        # core_j: (r_mid, n_j, r_right)
        core_i = self.cores[i]
        core_j = self.cores[i + 1]

        r_left, n_i, _ = core_i.shape
        _, n_j, r_right = core_j.shape

        # Contract into single tensor: (r_left, n_i, r_mid) @ (r_mid, n_j, r_right)
        # -> (r_left, n_i, n_j, r_right)
        contracted = torch.einsum('lim,mjr->lijr', core_i, core_j)

        # Swap physical dimensions: (r_left, n_i, n_j, r_right) -> (r_left, n_j, n_i, r_right)
        swapped = contracted.permute(0, 2, 1, 3)

        # Reshape for SVD: (r_left * n_j, n_i * r_right)
        mat = swapped.reshape(r_left * n_j, n_i * r_right)

        # SVD to re-split
        U, S, Vh = torch.linalg.svd(mat, full_matrices=False)

        # Apply truncation
        rank = S.shape[0]
        if max_rank is not None:
            rank = min(rank, max_rank)
        if rel_tol is not None:
            threshold = rel_tol * S[0]
            rank = min(rank, int((S > threshold).sum().item()))
        if abs_tol is not None:
            rank = min(rank, int((S > abs_tol).sum().item()))
        rank = max(rank, 1)  # Keep at least rank 1

        U = U[:, :rank]
        S = S[:rank]
        Vh = Vh[:rank, :]

        # Absorb singular values into U
        U = U @ torch.diag(S)

        # Reshape back to cores
        # new_core_i: (r_left, n_j, rank)
        # new_core_j: (rank, n_i, r_right)
        new_core_i = U.reshape(r_left, n_j, rank)
        new_core_j = Vh.reshape(rank, n_i, r_right)

        # Build new TensorTrain
        new_dims = self.dims[:i] + [n_j, n_i] + self.dims[i + 2:]
        new_ranks = self.ranks[1:i + 1] + [rank] + self.ranks[i + 2:-1]

        result = TensorTrain(new_dims, new_ranks if new_ranks else None,
                             self.dtype, self.device)

        # Copy cores before swap
        for k in range(i):
            result.cores[k] = self.cores[k].clone()

        # Set swapped cores
        result.cores[i] = new_core_i
        result.cores[i + 1] = new_core_j

        # Copy cores after swap
        for k in range(i + 2, self.d):
            result.cores[k] = self.cores[k].clone()

        return result

    def transpose(self, dim0: int, dim1: int) -> "TensorTrain":
        """
        Swap two modes (not necessarily adjacent).

        For non-adjacent modes, this is implemented as a sequence of
        adjacent swaps. The number of swaps is 2*|dim1-dim0| - 1.

        Args:
            dim0: First dimension to swap
            dim1: Second dimension to swap

        Returns:
            New TensorTrain with dimensions swapped

        Example:
            tt = TensorTrain.random([2, 3, 4, 5], ranks=[2, 3, 4])
            tt2 = tt.transpose(0, 2)  # shape becomes (4, 3, 2, 5)
        """
        if dim0 < 0 or dim0 >= self.d:
            raise ValueError(f"Invalid dim0={dim0} for TT with {self.d} modes")
        if dim1 < 0 or dim1 >= self.d:
            raise ValueError(f"Invalid dim1={dim1} for TT with {self.d} modes")

        if dim0 == dim1:
            return self.clone()

        # Ensure dim0 < dim1
        if dim0 > dim1:
            dim0, dim1 = dim1, dim0

        result = self.clone()

        # Step 1: Bubble element at dim1 left to position dim0
        # This takes (dim1 - dim0) swaps
        for j in range(dim1 - 1, dim0 - 1, -1):
            result = result.swap_adjacent(j)

        # After step 1: element originally at dim1 is now at dim0
        # Element originally at dim0 is now at dim0 + 1

        # Step 2: Bubble element originally at dim0 (now at dim0+1) right to dim1
        # This takes (dim1 - dim0 - 1) swaps
        for j in range(dim0 + 1, dim1):
            result = result.swap_adjacent(j)

        return result

    def unsqueeze(self, dim: int) -> "TensorTrain":
        """
        Insert a singleton dimension at the specified position.

        Adds a new mode of size 1 by inserting an identity-like core
        that preserves tensor values.

        Args:
            dim: Position where to insert the new dimension.
                 Supports negative indexing: -1 means after the last dim.

        Returns:
            New TensorTrain with an additional mode of size 1

        Example:
            tt = TensorTrain.random([2, 3, 4], ranks=[5, 6])
            tt2 = tt.unsqueeze(0)   # shape becomes (1, 2, 3, 4)
            tt3 = tt.unsqueeze(2)   # shape becomes (2, 3, 1, 4)
            tt4 = tt.unsqueeze(-1)  # shape becomes (2, 3, 4, 1)
        """
        # Handle negative indexing
        if dim < 0:
            dim = self.d + 1 + dim
        if dim < 0 or dim > self.d:
            raise ValueError(
                f"Invalid dim={dim} for unsqueeze on TT with {self.d} modes"
            )

        new_dims = self.dims[:dim] + [1] + self.dims[dim:]

        # Internal ranks are self.ranks[1:-1], constructor expects d-1 internal ranks
        # For new TT with d+1 dims, we need d internal ranks
        internal_ranks = self.ranks[1:-1]  # [r_1, ..., r_{d-1}]

        if dim == 0:
            # Insert at beginning: new bond between new core and old first core has rank 1
            # New internal ranks: [1, r_1, r_2, ..., r_{d-1}]
            new_internal_ranks = [1] + internal_ranks
            result = TensorTrain(new_dims, new_internal_ranks, self.dtype, self.device)

            # New core at position 0: (1, 1, 1) identity
            result.cores[0] = torch.ones(1, 1, 1, dtype=self.dtype, device=self.device)

            # Copy remaining cores
            for k in range(self.d):
                result.cores[k + 1] = self.cores[k].clone()

        elif dim == self.d:
            # Insert at end: new bond between old last core and new core has rank 1
            # New internal ranks: [r_1, r_2, ..., r_{d-1}, 1]
            new_internal_ranks = internal_ranks + [1]
            result = TensorTrain(new_dims, new_internal_ranks, self.dtype, self.device)

            # Copy existing cores
            for k in range(self.d):
                result.cores[k] = self.cores[k].clone()

            # New core at end: (1, 1, 1) identity
            result.cores[self.d] = torch.ones(1, 1, 1, dtype=self.dtype, device=self.device)

        else:
            # Insert in middle: identity core (r, 1, r) preserves the bond rank
            r = self.ranks[dim]  # rank of bond between core[dim-1] and core[dim]
            # internal_ranks[k] is the rank between core k and core k+1
            # We insert at position dim, so:
            # - bond (dim-1)-dim gets rank r (was internal_ranks[dim-1])
            # - bond dim-(dim+1) gets rank r (new identity core connects to old core dim)
            # - remaining bonds stay the same
            new_internal_ranks = internal_ranks[:dim-1] + [r, r] + internal_ranks[dim:]
            result = TensorTrain(new_dims, new_internal_ranks, self.dtype, self.device)

            # Copy cores before insertion point
            for k in range(dim):
                result.cores[k] = self.cores[k].clone()

            # New identity core: (r, 1, r) - acts as identity
            identity_core = torch.eye(r, dtype=self.dtype, device=self.device).unsqueeze(1)
            result.cores[dim] = identity_core

            # Copy cores after insertion point
            for k in range(dim, self.d):
                result.cores[k + 1] = self.cores[k].clone()

        return result

    def squeeze(self, dim: Optional[int] = None) -> "TensorTrain":
        """
        Remove singleton dimensions.

        If dim is specified, removes that dimension if it has size 1.
        If dim is None, removes all dimensions of size 1.

        Args:
            dim: Dimension to squeeze, or None to squeeze all singleton dims.
                 Supports negative indexing.

        Returns:
            New TensorTrain with singleton dimension(s) removed

        Raises:
            ValueError: If specified dim does not have size 1

        Example:
            tt = TensorTrain with shape (1, 3, 1, 4)
            tt.squeeze()    # shape becomes (3, 4)
            tt.squeeze(0)   # shape becomes (3, 1, 4)
            tt.squeeze(2)   # shape becomes (1, 3, 4)
        """
        if dim is not None:
            # Handle negative indexing
            if dim < 0:
                dim = self.d + dim
            if dim < 0 or dim >= self.d:
                raise ValueError(
                    f"Invalid dim={dim} for squeeze on TT with {self.d} modes"
                )
            if self.dims[dim] != 1:
                raise ValueError(
                    f"Cannot squeeze dim {dim}: size is {self.dims[dim]}, not 1"
                )

            # Squeeze single dimension
            return self._squeeze_dim(dim)

        else:
            # Squeeze all singleton dimensions
            result = self
            # Work backwards to avoid index shifting issues
            for k in range(self.d - 1, -1, -1):
                if result.dims[k] == 1:
                    result = result._squeeze_dim(k)
            return result

    def _squeeze_dim(self, dim: int) -> "TensorTrain":
        """
        Remove a single singleton dimension by contracting its core.

        Args:
            dim: Dimension to remove (must have size 1)

        Returns:
            New TensorTrain with dimension removed
        """
        if self.d == 1:
            raise ValueError("Cannot squeeze: would result in 0-dimensional TT")

        # Core to remove: (r_left, 1, r_right)
        core = self.cores[dim]

        # The core is essentially a matrix (r_left, r_right) after squeezing
        matrix = core.squeeze(1)  # (r_left, r_right)

        new_dims = self.dims[:dim] + self.dims[dim + 1:]

        if dim == 0:
            # Contract into next core
            # matrix: (1, r_right), next_core: (r_right, n, r_next)
            # Result: (1, n, r_next)
            next_core = self.cores[1]  # (r_right, n, r_next)
            new_first = torch.einsum('lr,rno->lno', matrix, next_core)

            new_ranks = self.ranks[2:-1]
            result = TensorTrain(new_dims, new_ranks, self.dtype, self.device)
            result.cores[0] = new_first
            for k in range(2, self.d):
                result.cores[k - 1] = self.cores[k].clone()

        elif dim == self.d - 1:
            # Contract into previous core
            # prev_core: (r_prev, n, r_left), matrix: (r_left, 1)
            # Result: (r_prev, n, 1)
            prev_core = self.cores[dim - 1]  # (r_prev, n, r_left)
            new_last = torch.einsum('pnl,lr->pnr', prev_core, matrix)

            new_ranks = self.ranks[1:-2]
            result = TensorTrain(new_dims, new_ranks, self.dtype, self.device)
            for k in range(dim - 1):
                result.cores[k] = self.cores[k].clone()
            result.cores[dim - 1] = new_last

        else:
            # Contract into next core (could also do previous, arbitrary choice)
            # matrix: (r_left, r_right), next_core: (r_right, n, r_next)
            # Result: (r_left, n, r_next)
            next_core = self.cores[dim + 1]  # (r_right, n, r_next)
            merged = torch.einsum('lr,rno->lno', matrix, next_core)

            # After removing core at dim and merging into dim+1:
            # - Bonds before dim stay the same: self.ranks[1:dim]
            # - Bond (dim-1)-dim becomes bond to merged core: self.ranks[dim]
            # - Bonds after dim+1 shift: self.ranks[dim+2:-1]
            new_internal_ranks = self.ranks[1:dim+1] + self.ranks[dim+2:-1]
            result = TensorTrain(new_dims, new_internal_ranks, self.dtype, self.device)

            for k in range(dim):
                result.cores[k] = self.cores[k].clone()
            result.cores[dim] = merged
            for k in range(dim + 2, self.d):
                result.cores[k - 1] = self.cores[k].clone()

        return result

    def to_tensor(self) -> torch.Tensor:
        """
        Contract the tensor train into the full tensor.

        Returns:
            Tensor of shape (n_1, n_2, ..., n_d)
        """
        # Start with the first core: (1, n_0, r_0) -> (n_0, r_0)
        T = self.cores[0].squeeze(0)

        # Contract with remaining cores
        for k in range(1, self.d):
            # T: (..., r_{k-1}), core[k]: (r_{k-1}, n_k, r_k)
            # Result: (..., n_k, r_k)
            T = torch.einsum('...i,ijk->...jk', T, self.cores[k])

        # Remove final rank-1 dimension: (..., 1) -> (...)
        return T.squeeze(-1)

    def to_dense(self) -> torch.Tensor:
        """
        Contract the tensor train into the full dense tensor (torch-like API).

        Returns:
            Tensor of shape (n_1, n_2, ..., n_d)
        """
        return self.to_tensor()


def _compute_reshape_ops(old_dims: List[int], new_dims: List[int]) -> List[tuple]:
    """
    Compute the sequence of merge/split operations to reshape old_dims to new_dims.

    Returns a list of operations: ('merge', (start, end)) or ('split', (mode, shape))

    Raises ValueError if the reshape requires permutation.
    """
    ops = []
    old_idx = 0
    new_idx = 0
    current_dims = list(old_dims)

    while new_idx < len(new_dims):
        if old_idx >= len(current_dims):
            raise ValueError(
                f"Cannot reshape {old_dims} to {new_dims}: incompatible structure"
            )

        old_prod = current_dims[old_idx]
        new_target = new_dims[new_idx]

        if old_prod == new_target:
            # Dimensions match, move forward
            old_idx += 1
            new_idx += 1

        elif old_prod < new_target:
            # Need to merge: accumulate old dims until we match or exceed
            merge_end = old_idx + 1
            while old_prod < new_target and merge_end < len(current_dims):
                old_prod *= current_dims[merge_end]
                merge_end += 1

            if old_prod != new_target:
                raise ValueError(
                    f"Cannot reshape {old_dims} to {new_dims}: "
                    f"no valid merge found at position {old_idx}"
                )

            # Record merge operation (on current state, not original)
            # We need to track the offset due to previous operations
            ops.append(('merge', (old_idx, merge_end)))

            # Update current_dims to reflect the merge
            merged_dim = math.prod(current_dims[old_idx:merge_end])
            current_dims = current_dims[:old_idx] + [merged_dim] + current_dims[merge_end:]

            old_idx += 1
            new_idx += 1

        else:  # old_prod > new_target
            # Need to split: find shape from new_dims that multiply to old_prod
            split_end = new_idx + 1
            new_prod = new_target
            while new_prod < old_prod and split_end < len(new_dims):
                new_prod *= new_dims[split_end]
                split_end += 1

            if new_prod != old_prod:
                raise ValueError(
                    f"Cannot reshape {old_dims} to {new_dims}: "
                    f"no valid split found at position {old_idx}"
                )

            split_shape = new_dims[new_idx:split_end]
            ops.append(('split', (old_idx, split_shape)))

            # Update current_dims to reflect the split
            current_dims = current_dims[:old_idx] + split_shape + current_dims[old_idx+1:]

            old_idx += len(split_shape)
            new_idx = split_end

    if old_idx < len(current_dims):
        raise ValueError(
            f"Cannot reshape {old_dims} to {new_dims}: leftover dimensions"
        )

    return ops


def dot(tt1: TensorTrain, tt2: TensorTrain) -> List[torch.Tensor]:
    """
    Compute the tensor train dot product via left-to-right sweep.

    Performs efficient O(n·r³) contraction by maintaining intermediate
    boundary matrices L_k of shape (r1_k, r2_k) during the sweep.

    The contraction at each site:
        L_k[j, l] = sum_{i,m,n} L_{k-1}[i,m] * A[k][i,n,j] * B[k][m,n,l]

    Args:
        tt1: First tensor train
        tt2: Second tensor train

    Returns:
        List of d+1 boundary matrices L_0, L_1, ..., L_d where
        L_k has shape (r1_k, r2_k). L_0 is (1,1) with value 1,
        and L_d is (1,1) containing the inner product.
    """
    if tt1.d != tt2.d:
        raise ValueError(
            f"TTs must have same number of cores: {tt1.d} vs {tt2.d}"
        )
    if tt1.dims != tt2.dims:
        raise ValueError(
            f"TTs must have same dimensions: {tt1.dims} vs {tt2.dims}"
        )

    # Initialize L_0 as 1x1 identity (r_0 = 1 for both TTs)
    L = torch.ones(1, 1, dtype=tt1.dtype, device=tt1.device)
    matrices: List[torch.Tensor] = [L]

    for k in range(tt1.d):
        A = tt1.cores[k]  # (rA_left, n, rA_right)
        B = tt2.cores[k]  # (rB_left, n, rB_right)

        # Step 1: contract L with A over left index of A
        # L[i,m], A[i,n,j] -> temp[m,n,j]
        temp = torch.einsum('im,inj->mnj', L, A)

        # Step 2: contract temp with B over left index of B and physical index
        # temp[m,n,j], B[m,n,l] -> L_new[j,l]
        L = torch.einsum('mnj,mnl->jl', temp, B)

        matrices.append(L)

    return matrices


def hadamard(tt1: TensorTrain, tt2: TensorTrain) -> TensorTrain:
    """
    Compute the Hadamard (element-wise) product of two tensor trains.

    The result has ranks that are products of the input ranks:
    r_result_k = r1_k * r2_k

    Args:
        tt1: First tensor train
        tt2: Second tensor train

    Returns:
        TensorTrain representing the element-wise product
    """
    if tt1.d != tt2.d:
        raise ValueError(
            f"TTs must have same number of cores: {tt1.d} vs {tt2.d}"
        )
    if tt1.dims != tt2.dims:
        raise ValueError(
            f"TTs must have same dimensions: {tt1.dims} vs {tt2.dims}"
        )

    # Result ranks are products of input ranks
    result_ranks = [tt1.ranks[k] * tt2.ranks[k] for k in range(1, tt1.d)]
    result = TensorTrain(tt1.dims, result_ranks, tt1.dtype, tt1.device)

    for k in range(tt1.d):
        A = tt1.cores[k]  # (rA_left, n, rA_right)
        B = tt2.cores[k]  # (rB_left, n, rB_right)

        rA_left, n, rA_right = A.shape
        rB_left, _, rB_right = B.shape

        # C[(i,i'), n, (j,j')] = A[i,n,j] * B[i',n,j']
        # einsum: 'inj,knl->iknjl' then reshape
        C = torch.einsum('inj,knl->iknjl', A, B)
        C = C.reshape(rA_left * rB_left, n, rA_right * rB_right)

        result.cores[k] = C

    return result


def scale(tt: TensorTrain, alpha: float) -> TensorTrain:
    """
    Multiply a tensor train by a scalar.

    Args:
        tt: Tensor train
        alpha: Scalar multiplier

    Returns:
        New TensorTrain representing alpha * tt
    """
    result = tt.copy()
    result.cores[0] = alpha * result.cores[0]
    return result


def add(tt1: TensorTrain, tt2: TensorTrain) -> TensorTrain:
    """
    Add two tensor trains.

    The result has ranks that are sums of the input ranks:
    r_result_k = r1_k + r2_k (for internal ranks)

    Args:
        tt1: First tensor train
        tt2: Second tensor train

    Returns:
        TensorTrain representing tt1 + tt2
    """
    if tt1.d != tt2.d:
        raise ValueError(
            f"TTs must have same number of cores: {tt1.d} vs {tt2.d}"
        )
    if tt1.dims != tt2.dims:
        raise ValueError(
            f"TTs must have same dimensions: {tt1.dims} vs {tt2.dims}"
        )

    # Result ranks are sums of input ranks (except boundaries which stay 1)
    result_ranks = [tt1.ranks[k] + tt2.ranks[k] for k in range(1, tt1.d)]
    result = TensorTrain(tt1.dims, result_ranks, tt1.dtype, tt1.device)

    for k in range(tt1.d):
        A = tt1.cores[k]  # (rA_left, n, rA_right)
        B = tt2.cores[k]  # (rB_left, n, rB_right)

        rA_left, n, rA_right = A.shape
        rB_left, _, rB_right = B.shape

        if k == 0:
            # First core: concatenate along right rank
            # (1, n, rA) and (1, n, rB) -> (1, n, rA + rB)
            C = torch.cat([A, B], dim=2)
        elif k == tt1.d - 1:
            # Last core: concatenate along left rank
            # (rA, n, 1) and (rB, n, 1) -> (rA + rB, n, 1)
            C = torch.cat([A, B], dim=0)
        else:
            # Middle cores: block diagonal
            # Top-left: A, Bottom-right: B, rest zeros
            C = torch.zeros(
                rA_left + rB_left, n, rA_right + rB_right,
                dtype=tt1.dtype, device=tt1.device
            )
            C[:rA_left, :, :rA_right] = A
            C[rA_left:, :, rA_right:] = B

        result.cores[k] = C

    return result


def cat(tensors: List[TensorTrain], dim: int = 0) -> TensorTrain:
    """
    Concatenate tensor trains along an existing dimension.

    All tensor trains must have the same shape except at the concatenation
    dimension. The result has ranks that are sums of input ranks.

    Args:
        tensors: List of TensorTrains to concatenate
        dim: Dimension along which to concatenate (supports negative indexing)

    Returns:
        TensorTrain with concatenated dimension

    Example:
        tt1 = TensorTrain with shape (2, 3, 4)
        tt2 = TensorTrain with shape (5, 3, 4)
        cat([tt1, tt2], dim=0)  # shape (7, 3, 4)
    """
    if len(tensors) == 0:
        raise ValueError("Need at least one tensor to concatenate")
    if len(tensors) == 1:
        return tensors[0].clone()

    # Reduce pairwise
    result = tensors[0]
    for tt in tensors[1:]:
        result = _cat_two(result, tt, dim)
    return result


def _cat_two(tt1: TensorTrain, tt2: TensorTrain, dim: int) -> TensorTrain:
    """Concatenate two TensorTrains along dimension dim."""
    # Handle negative indexing
    if dim < 0:
        dim = tt1.d + dim

    if dim < 0 or dim >= tt1.d:
        raise ValueError(f"Invalid dim={dim} for TT with {tt1.d} modes")

    if tt1.d != tt2.d:
        raise ValueError(
            f"TTs must have same number of modes: {tt1.d} vs {tt2.d}"
        )

    # Check shapes match except at dim
    for k in range(tt1.d):
        if k != dim and tt1.dims[k] != tt2.dims[k]:
            raise ValueError(
                f"Dimension mismatch at mode {k}: {tt1.dims[k]} vs {tt2.dims[k]}. "
                f"All dimensions except dim={dim} must match."
            )

    # Build result dims and ranks
    result_dims = list(tt1.dims)
    result_dims[dim] = tt1.dims[dim] + tt2.dims[dim]

    # Result internal ranks are sums (except boundaries which stay 1)
    result_internal_ranks = [
        tt1.ranks[k] + tt2.ranks[k] for k in range(1, tt1.d)
    ]

    result = TensorTrain(result_dims, result_internal_ranks, tt1.dtype, tt1.device)

    for k in range(tt1.d):
        A = tt1.cores[k]  # (rA_left, n1, rA_right)
        B = tt2.cores[k]  # (rB_left, n2, rB_right)

        rA_left, n1, rA_right = A.shape
        rB_left, n2, rB_right = B.shape

        if k == dim:
            # Concatenation dimension: stack physical indices with block structure
            # Result shape: (rA_left + rB_left, n1 + n2, rA_right + rB_right)
            # But handle boundary cases

            if k == 0:
                # First core: left rank is 1 for both
                # (1, n1, rA) cat (1, n2, rB) -> (1, n1+n2, rA+rB)
                C = torch.zeros(
                    1, n1 + n2, rA_right + rB_right,
                    dtype=tt1.dtype, device=tt1.device
                )
                C[0, :n1, :rA_right] = A[0, :, :]
                C[0, n1:, rA_right:] = B[0, :, :]

            elif k == tt1.d - 1:
                # Last core: right rank is 1 for both
                # (rA, n1, 1) cat (rB, n2, 1) -> (rA+rB, n1+n2, 1)
                C = torch.zeros(
                    rA_left + rB_left, n1 + n2, 1,
                    dtype=tt1.dtype, device=tt1.device
                )
                C[:rA_left, :n1, 0] = A[:, :, 0]
                C[rA_left:, n1:, 0] = B[:, :, 0]

            else:
                # Middle core at concatenation dim
                C = torch.zeros(
                    rA_left + rB_left, n1 + n2, rA_right + rB_right,
                    dtype=tt1.dtype, device=tt1.device
                )
                C[:rA_left, :n1, :rA_right] = A
                C[rA_left:, n1:, rA_right:] = B

        else:
            # Non-concatenation dimension: block diagonal structure
            # Physical dimensions must match (already validated)
            n = n1  # n1 == n2

            if k == 0:
                # First core: (1, n, rA+rB)
                C = torch.cat([A, B], dim=2)

            elif k == tt1.d - 1:
                # Last core: (rA+rB, n, 1)
                C = torch.cat([A, B], dim=0)

            else:
                # Middle core: block diagonal
                C = torch.zeros(
                    rA_left + rB_left, n, rA_right + rB_right,
                    dtype=tt1.dtype, device=tt1.device
                )
                C[:rA_left, :, :rA_right] = A
                C[rA_left:, :, rA_right:] = B

        result.cores[k] = C

    return result


def stack(tensors: List[TensorTrain], dim: int = 0) -> TensorTrain:
    """
    Stack tensor trains along a new dimension.

    All tensor trains must have exactly the same shape. A new dimension
    of size len(tensors) is inserted at position dim.

    Args:
        tensors: List of TensorTrains to stack
        dim: Position for new dimension (supports negative indexing)

    Returns:
        TensorTrain with new dimension

    Example:
        tt1 = TensorTrain with shape (2, 3, 4)
        tt2 = TensorTrain with shape (2, 3, 4)
        stack([tt1, tt2], dim=0)  # shape (2, 2, 3, 4)
        stack([tt1, tt2], dim=1)  # shape (2, 2, 3, 4)
    """
    if len(tensors) == 0:
        raise ValueError("Need at least one tensor to stack")

    # Validate all shapes are identical
    first_shape = tensors[0].shape
    for i, tt in enumerate(tensors[1:], 1):
        if tt.shape != first_shape:
            raise ValueError(
                f"All tensors must have same shape. "
                f"Tensor 0 has shape {first_shape}, tensor {i} has shape {tt.shape}"
            )

    # Unsqueeze all tensors at the stack dimension, then concatenate
    expanded = [tt.unsqueeze(dim) for tt in tensors]
    return cat(expanded, dim=dim)
