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


class TTMatrix:
    """
    Tensor train representation of a linear operator (TT-matrix / MPO).

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

    The TTMatrix represents a linear operator mapping vectors of dimension
    (n_col_1 * n_col_2 * ... * n_col_d) to vectors of dimension
    (n_row_1 * n_row_2 * ... * n_row_d).
    """

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

        Args:
            dims: List of (n_row, n_col) tuples for each site
            ranks: List of internal ranks [r_1, r_2, ..., r_{d-1}].
                   If None, defaults to all ones (rank-1 TTMatrix).
            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_row_k, n_col_k, r_k)
        self.cores: List[torch.Tensor] = []
        for k in range(self.d):
            n_row, n_col = self._dims[k]
            shape = (self.ranks[k], n_row, n_col, self.ranks[k + 1])
            core = torch.zeros(shape, dtype=dtype, device=device)
            self.cores.append(core)

    @property
    def row_dims(self) -> List[int]:
        """Return the row dimensions [n_row_1, n_row_2, ..., n_row_d]."""
        return [d[0] for d in self._dims]

    @property
    def col_dims(self) -> List[int]:
        """Return the column dimensions [n_col_1, n_col_2, ..., n_col_d]."""
        return [d[1] for d in self._dims]

    @property
    def dims(self) -> List[Tuple[int, int]]:
        """Return the list of (n_row, n_col) tuples."""
        return self._dims

    @property
    def shape(self) -> Tuple[int, int]:
        """Return the shape as a linear operator (total_rows, total_cols)."""
        total_rows = 1
        total_cols = 1
        for n_row, n_col in self._dims:
            total_rows *= n_row
            total_cols *= n_col
        return (total_rows, total_cols)

    @property
    def ndim(self) -> int:
        """Return the number of sites (degrees of freedom)."""
        return self.d

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

    @property
    def numel(self) -> int:
        """Logical number of matrix elements: total_rows * total_cols."""
        total_rows, total_cols = self.shape
        return total_rows * total_cols

    @property
    def ttnumel(self) -> int:
        """Actual storage size: sum of r_{k-1} * n_row_k * n_col_k * r_k."""
        return sum(
            self.ranks[k] * self._dims[k][0] * self._dims[k][1] * self.ranks[k + 1]
            for k in range(self.d)
        )

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

    @classmethod
    def identity(
        cls,
        dims: List[int],
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ) -> "TTMatrix":
        """
        Create an identity TT-matrix.

        Args:
            dims: List of dimensions [n_1, n_2, ..., n_d] (square at each site)
            dtype: Torch dtype
            device: Torch device

        Returns:
            Rank-1 TTMatrix representing the identity operator
        """
        # Identity has square dims at each site
        square_dims = [(n, n) for n in dims]
        ttm = cls(square_dims, None, dtype, device)  # rank-1 by default

        for k in range(ttm.d):
            n = dims[k]
            # Core[k] should be identity: core[0, i, j, 0] = delta_{ij}
            ttm.cores[k] = torch.eye(n, dtype=dtype, device=device).reshape(1, n, n, 1)

        return ttm

    @classmethod
    def zeros(
        cls,
        dims: List[Tuple[int, int]],
        ranks: Optional[List[int]] = None,
        dtype: torch.dtype = torch.float64,
        device: Optional[torch.device] = None,
    ) -> "TTMatrix":
        """Create a TT-matrix with all zero entries."""
        # Already initialized to zeros by default
        return cls(dims, ranks, dtype, device)

    @classmethod
    def local(
        cls,
        matrix: torch.Tensor,
    ) -> "TTMatrix":
        """
        Create a single-site TTMatrix from a dense matrix.

        This wraps a 2D tensor as a 1-site TTMatrix without any decomposition.
        Useful for creating local operators that can be applied to specific
        sites using the .on() syntax.

        Args:
            matrix: 2D tensor of shape (n_row, n_col)

        Returns:
            1-site TTMatrix with core shape (1, n_row, n_col, 1)

        Example:
            # Create a local operator and apply to site 2
            X = torch.tensor([[0., 1.], [1., 0.]])
            flip = TTMatrix.local(X)
            result = flip.on(2) @ psi

            # Works with rectangular matrices too
            proj = TTMatrix.local(torch.randn(3, 5))  # Maps dim 5 -> 3
        """
        if matrix.ndim != 2:
            raise ValueError(f"Expected 2D matrix, got {matrix.ndim}D tensor")

        n_row, n_col = matrix.shape
        ttm = cls([(n_row, n_col)], ranks=[], dtype=matrix.dtype, device=matrix.device)
        ttm.cores[0] = matrix.reshape(1, n_row, n_col, 1)
        return ttm

    def __getitem__(self, k: int) -> torch.Tensor:
        """Get the k-th core tensor."""
        return self.cores[k]

    def __setitem__(self, k: int, value: torch.Tensor):
        """Set the k-th core tensor."""
        n_row, n_col = self._dims[k]
        expected_shape = (self.ranks[k], n_row, n_col, 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 (sites)."""
        return self.d

    def copy(self) -> "TTMatrix":
        """Return a deep copy of this TT-matrix."""
        ttm = TTMatrix(self._dims, self.ranks[1:-1], self.dtype, self.device)
        for k in range(self.d):
            ttm.cores[k] = self.cores[k].clone()
        return ttm

    def clone(self) -> "TTMatrix":
        """Return a deep copy of this TT-matrix (torch-like API)."""
        return self.copy()

    def detach(self) -> "TTMatrix":
        """
        Detach TT-matrix from computation graph (torch-like API).

        Returns a new TTMatrix with all cores detached from the
        autograd computation graph.

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

    def to(
        self,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
    ) -> "TTMatrix":
        """
        Convert TT-matrix 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 TTMatrix with converted cores
        """
        if dtype is None:
            dtype = self.dtype
        if device is None:
            device = self.device

        ttm = TTMatrix(self._dims, self.ranks[1:-1], dtype, device)
        for k in range(self.d):
            ttm.cores[k] = self.cores[k].to(dtype=dtype, device=device)
        return ttm

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

        Returns:
            List of tuples [(r0, n_row_1, n_col_1, r1), ...]
        """
        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 tt_cores(self) -> List[torch.Tensor]:
        """
        Return the list of TT-matrix cores.

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

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

    def total_params(self) -> int:
        """Return total number of parameters in the TT-matrix."""
        return sum(core.numel() for core in 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 matrix size to stored tensor size
        """
        return self.numel / self.storage_numel()

    def to_tensor(self) -> torch.Tensor:
        """
        Contract the TT-matrix into the full dense matrix.

        Returns:
            Tensor of shape (n_row_1 * ... * n_row_d, n_col_1 * ... * n_col_d)
        """
        # Start with the first core: (1, n_row, n_col, r) -> (n_row, n_col, r)
        T = self.cores[0].squeeze(0)

        # Contract with remaining cores
        for k in range(1, self.d):
            # T: (..., n_row_prev, n_col_prev, r), core[k]: (r, n_row, n_col, r_next)
            # Result: (..., n_row_prev, n_col_prev, n_row, n_col, r_next)
            T = torch.einsum('...i,ijkl->...jkl', T, self.cores[k])

        # T now has shape (n_row_1, n_col_1, n_row_2, n_col_2, ..., n_row_d, n_col_d, 1)
        T = T.squeeze(-1)

        # Reshape to matrix form
        # Current shape: (n_row_1, n_col_1, n_row_2, n_col_2, ..., n_row_d, n_col_d)
        # Need to permute to: (n_row_1, n_row_2, ..., n_row_d, n_col_1, n_col_2, ..., n_col_d)
        # Then reshape to (total_rows, total_cols)

        # Build permutation: [0, 2, 4, ..., 1, 3, 5, ...]
        row_indices = list(range(0, 2 * self.d, 2))
        col_indices = list(range(1, 2 * self.d, 2))
        perm = row_indices + col_indices

        T = T.permute(*perm)

        # Reshape to matrix
        total_rows, total_cols = self.shape
        return T.reshape(total_rows, total_cols)

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

        Returns:
            Tensor of shape (n_row_1 * ... * n_row_d, n_col_1 * ... * n_col_d)
        """
        return self.to_tensor()

    def to_tensor_train(self) -> "TensorTrain":
        """
        Convert TTMatrix to TensorTrain by merging row and column dimensions.

        Each 4D core (r_l, n_row, n_col, r_r) becomes a 3D core (r_l, n_row*n_col, r_r).
        The resulting TensorTrain has dims [n_row_1*n_col_1, n_row_2*n_col_2, ...].

        Returns:
            TensorTrain with merged physical dimensions

        Example:
            ttm = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])
            tt = ttm.to_tensor_train()  # dims = [12, 30]
        """
        # Import here to avoid circular import at module load
        from .tensortrain import TensorTrain

        # Merged dims: n_row * n_col at each site
        merged_dims = [n_row * n_col for n_row, n_col in self._dims]

        # Create TensorTrain with same ranks
        tt = TensorTrain(merged_dims, self.ranks[1:-1], self.dtype, self.device)

        # Reshape each core: (r_l, n_row, n_col, r_r) -> (r_l, n_row*n_col, r_r)
        for k in range(self.d):
            r_l, n_row, n_col, r_r = self.cores[k].shape
            tt.cores[k] = self.cores[k].reshape(r_l, n_row * n_col, r_r)

        return tt

    def __matmul__(self, other):
        """
        Matrix multiplication: self @ other.

        If other is a TTMatrix, computes matmul(self, other).
        If other is a TensorTrain, computes apply(self, other).
        """
        # Import here to avoid issues at class definition time
        from .tensortrain import TensorTrain

        if isinstance(other, TTMatrix):
            return matmul(self, other)
        elif isinstance(other, TensorTrain):
            return apply(self, other)
        else:
            return NotImplemented

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

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

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

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

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

    def __and__(self, other: "TTMatrix") -> "TTMatrix":
        """Kronecker product: self & other."""
        if isinstance(other, TTMatrix):
            return kron(self, other)
        return NotImplemented

    def on(self, *sites: int) -> "LocatedOp":
        """
        Specify which sites this operator acts on.

        Returns a LocatedOp that can be applied to a TensorTrain using @.

        Args:
            *sites: Site indices this operator acts on

        Returns:
            LocatedOp wrapper that applies this operator to specified sites

        Example:
            # 2-site gate acting on sites 1 and 2 of a 5-site state
            H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
            psi = TensorTrain.random([2, 3, 4, 5, 6], ranks=[2, 2, 2, 2])
            result = H.on(1, 2) @ psi
        """
        return LocatedOp(self, sites)

    @classmethod
    def from_tensor_train(
        cls,
        tt: "TensorTrain",
        dims: List[Tuple[int, int]],
    ) -> "TTMatrix":
        """
        Create TTMatrix from TensorTrain by splitting physical dimensions.

        Each 3D core (r_l, n, r_r) becomes a 4D core (r_l, n_row, n_col, r_r)
        where n = n_row * n_col.

        Args:
            tt: TensorTrain to convert
            dims: List of (n_row, n_col) tuples specifying how to split each dimension

        Returns:
            TTMatrix with split physical dimensions

        Raises:
            ValueError: If dims don't match TensorTrain dimensions

        Example:
            tt = TensorTrain.random([12, 30], ranks=[2])
            ttm = TTMatrix.from_tensor_train(tt, dims=[(3, 4), (5, 6)])
        """
        if len(dims) != tt.d:
            raise ValueError(
                f"Number of dims ({len(dims)}) must match TensorTrain sites ({tt.d})"
            )

        for k, (n_row, n_col) in enumerate(dims):
            if n_row * n_col != tt.dims[k]:
                raise ValueError(
                    f"dims[{k}] = ({n_row}, {n_col}) product {n_row * n_col} "
                    f"!= TensorTrain dim[{k}] = {tt.dims[k]}"
                )

        # Create TTMatrix with same ranks
        ttm = cls(dims, tt.ranks[1:-1], tt.dtype, tt.device)

        # Reshape each core: (r_l, n_row*n_col, r_r) -> (r_l, n_row, n_col, r_r)
        for k in range(tt.d):
            r_l, _, r_r = tt.cores[k].shape
            n_row, n_col = dims[k]
            ttm.cores[k] = tt.cores[k].reshape(r_l, n_row, n_col, r_r)

        return ttm


# ============================================================================
# LocatedOp: TTMatrix with specified target sites
# ============================================================================


class LocatedOp:
    """
    A TTMatrix operator annotated with the sites it acts on.

    This wrapper enables natural syntax for partial application:
        result = H.on(1, 2) @ psi

    instead of:
        result = apply(H, psi, dims=[1, 2])
    """

    def __init__(self, op: TTMatrix, sites: Tuple[int, ...]):
        """
        Create a located operator.

        Args:
            op: The TTMatrix operator
            sites: Tuple of site indices this operator acts on
        """
        self.op = op
        self.sites = tuple(sites)

        if len(sites) != op.d:
            raise ValueError(
                f"Number of sites ({len(sites)}) must match operator sites ({op.d})"
            )

    def __matmul__(self, psi: "TensorTrain") -> "TensorTrain":
        """Apply this operator to a TensorTrain at the specified sites."""
        return apply(self.op, psi, dims=list(self.sites))

    def __repr__(self) -> str:
        return f"LocatedOp({self.op}, sites={self.sites})"


# ============================================================================
# Trunc: Truncation operator for circuit-style syntax
# ============================================================================


class Trunc:
    """
    Truncation operator for use in circuit-style expressions.

    This enables natural syntax for interleaving operators with truncation:
        result = Trunc(10) @ U.on(2, 3) @ Trunc(10) @ U.on(1, 2) @ psi

    The expression is evaluated right-to-left:
    1. Apply U to sites 1, 2
    2. Truncate to rank 10
    3. Apply U to sites 2, 3
    4. Truncate to rank 10
    """

    def __init__(
        self,
        max_rank: Optional[int] = None,
        rel_tol: Optional[float] = None,
        abs_tol: Optional[float] = None,
    ):
        """
        Create a truncation operator.

        Args:
            max_rank: Maximum rank to keep at each bond
            rel_tol: Relative tolerance for singular values
            abs_tol: Absolute tolerance for singular values

        Example:
            # Rank-based truncation
            result = Trunc(10) @ U.on(1, 2) @ psi

            # Tolerance-based truncation
            result = Trunc(rel_tol=1e-7) @ U.on(1, 2) @ psi
        """
        self.max_rank = max_rank
        self.rel_tol = rel_tol
        self.abs_tol = abs_tol

    def __matmul__(self, other):
        """
        Apply truncation: Trunc(...) @ tt.

        Works with TensorTrain and TTMatrix.
        """
        # Import here to avoid circular imports
        from . import linalg

        # Check if it's a TensorTrain
        from .tensortrain import TensorTrain
        if isinstance(other, TensorTrain):
            return linalg.svd(
                other,
                max_rank=self.max_rank,
                rel_tol=self.rel_tol,
                abs_tol=self.abs_tol,
            )

        # Check if it's a TTMatrix
        if isinstance(other, TTMatrix):
            return round(
                other,
                max_rank=self.max_rank,
                rel_tol=self.rel_tol,
                abs_tol=self.abs_tol,
            )

        return NotImplemented

    def __repr__(self) -> str:
        parts = []
        if self.max_rank is not None:
            parts.append(f"max_rank={self.max_rank}")
        if self.rel_tol is not None:
            parts.append(f"rel_tol={self.rel_tol}")
        if self.abs_tol is not None:
            parts.append(f"abs_tol={self.abs_tol}")
        return f"Trunc({', '.join(parts)})"


# ============================================================================
# Circuit: Composable chain of operators for circuit-style syntax
# ============================================================================


class Circuit:
    """
    A composable chain of operators for circuit-style expressions.

    Circuit enables natural syntax for interleaving operators with truncation:
        result = Trunc(10) @ U.on(2, 3) @ Trunc(10) @ U.on(1, 2) @ psi

    The expression builds a Circuit from left-to-right, then applies
    all operations right-to-left when a TensorTrain is encountered.

    This class is typically not constructed directly, but is created
    automatically when composing Trunc and LocatedOp objects.
    """

    def __init__(self, ops: List = None):
        """
        Create a circuit from a list of operations.

        Args:
            ops: List of operations (Trunc, LocatedOp, TTMatrix)
                 Operations are stored in left-to-right order but
                 applied in reverse (right-to-left) order.
        """
        self.ops = ops if ops is not None else []

    def __matmul__(self, other):
        """
        Extend the circuit or apply to TensorTrain.

        Trunc @ LocatedOp @ Trunc @ LocatedOp @ psi
        """
        from .tensortrain import TensorTrain

        if isinstance(other, TensorTrain):
            # Apply all ops in reverse order (right-to-left)
            result = other
            for op in reversed(self.ops):
                result = op @ result
            return result
        elif isinstance(other, (LocatedOp, Trunc, TTMatrix)):
            # Extend the circuit
            return Circuit(self.ops + [other])
        elif isinstance(other, Circuit):
            # Merge circuits
            return Circuit(self.ops + other.ops)
        return NotImplemented

    def __repr__(self) -> str:
        return f"Circuit({self.ops})"


# Update Trunc to create Circuit when combined with LocatedOp
_original_trunc_matmul = Trunc.__matmul__


def _new_trunc_matmul(self, other):
    """Extended __matmul__ that handles LocatedOp via Circuit."""
    if isinstance(other, LocatedOp):
        # Create a Circuit with Trunc first (applied last), then LocatedOp
        return Circuit([self, other])
    elif isinstance(other, Circuit):
        # Prepend Trunc to the circuit
        return Circuit([self] + other.ops)
    return _original_trunc_matmul(self, other)


Trunc.__matmul__ = _new_trunc_matmul


# ============================================================================
# TTMatrix Operations
# ============================================================================

# Import TensorTrain here to avoid circular imports at module load time
from .tensortrain import TensorTrain


def apply(
    ttm: TTMatrix,
    tt: TensorTrain,
    dims: Optional[Union[List[int], Tuple[List[int], List[int]]]] = None,
) -> TensorTrain:
    """
    Apply a TTMatrix operator to a TensorTrain vector.

    Computes result = ttm @ tt, where ttm acts on specified dimensions of tt.

    Args:
        ttm: TTMatrix operator with d_m sites
        tt: TensorTrain vector with d_v sites
        dims: Specification of which dimensions ttm acts on:
            - None: ttm acts on dims [0, 1, ..., d_m-1] of tt
            - List[int]: ttm acts on these dims of tt (len must equal d_m)
            - Tuple of two lists ([ttm_dims], [tt_dims]): explicit pairing

    Returns:
        TensorTrain with ttm.row_dims at the specified dimensions

    Example:
        # Full apply (H acts on all of psi)
        result = apply(H, psi)

        # Partial apply (2-site gate on dims 2,3 of a 5-site state)
        result = apply(gate, psi, dims=[2, 3])

        # Explicit pairing (H's dim 0 acts on psi's dim 2)
        result = apply(H, psi, dims=([0, 1], [2, 0]))
    """
    # Parse dims argument
    if dims is None:
        # Default: apply to first d_m dims of tt
        dims_ttm = list(range(ttm.d))
        dims_tt = list(range(ttm.d))
    elif isinstance(dims, tuple) and len(dims) == 2:
        dims_ttm, dims_tt = dims
        dims_ttm = list(dims_ttm)
        dims_tt = list(dims_tt)
    else:
        # dims is a list: apply ttm (in order) to these dims of tt
        dims_ttm = list(range(ttm.d))
        dims_tt = list(dims)

    # Validate
    if len(dims_ttm) != len(dims_tt):
        raise ValueError(
            f"dims_ttm and dims_tt must have same length: {len(dims_ttm)} vs {len(dims_tt)}"
        )
    if len(dims_tt) != ttm.d:
        raise ValueError(
            f"Number of dims ({len(dims_tt)}) must match TTMatrix sites ({ttm.d})"
        )

    for i, d in enumerate(dims_tt):
        if d < 0 or d >= tt.d:
            raise ValueError(f"dims[{i}]={d} out of range for TensorTrain with {tt.d} dims")
        ttm_idx = dims_ttm[i]
        if ttm.col_dims[ttm_idx] != tt.dims[d]:
            raise ValueError(
                f"Dimension mismatch: TTMatrix col_dim[{ttm_idx}]={ttm.col_dims[ttm_idx]} "
                f"!= TensorTrain dim[{d}]={tt.dims[d]}"
            )

    # Check if dims_tt is sorted (contiguous application is simpler)
    sorted_dims_tt = sorted(dims_tt)
    is_contiguous = (sorted_dims_tt == list(range(sorted_dims_tt[0], sorted_dims_tt[0] + len(sorted_dims_tt))))
    is_ordered = (dims_tt == sorted_dims_tt)

    if is_contiguous and is_ordered:
        return _apply_contiguous(ttm, tt, dims_tt)
    else:
        return _apply_general(ttm, tt, dims_ttm, dims_tt)


def _apply_contiguous(
    ttm: TTMatrix,
    tt: TensorTrain,
    dims_tt: List[int],
) -> TensorTrain:
    """
    Apply TTMatrix to contiguous, ordered dimensions of TensorTrain.

    This is the efficient case where dims_tt = [start, start+1, ..., start+k-1].
    """
    start = dims_tt[0]
    end = dims_tt[-1] + 1  # exclusive

    # Build result dimensions
    result_dims = list(tt.dims)
    for i, d in enumerate(dims_tt):
        result_dims[d] = ttm.row_dims[i]

    # Build result ranks
    # Before active region: tt ranks
    # In active region: tt × ttm ranks
    # After active region: tt ranks
    result_ranks = []
    for k in range(tt.d - 1):
        if k < start:
            result_ranks.append(tt.ranks[k + 1])
        elif k < end - 1:
            # Inside active region: product of ranks
            ttm_idx = k - start
            result_ranks.append(tt.ranks[k + 1] * ttm.ranks[ttm_idx + 1])
        elif k == end - 1:
            # Exiting active region: ttm rank goes to 1
            result_ranks.append(tt.ranks[k + 1])
        else:
            result_ranks.append(tt.ranks[k + 1])

    result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                         tt.dtype, tt.device)

    # Build result cores
    for k in range(tt.d):
        if k < start:
            # Before active region: copy tt core
            result.cores[k] = tt.cores[k].clone()
        elif k < end:
            # In active region: contract ttm with tt
            ttm_idx = k - start
            h_core = ttm.cores[ttm_idx]  # (r_h_l, n_row, n_col, r_h_r)
            psi_core = tt.cores[k]       # (r_psi_l, n_col, r_psi_r)

            r_h_l, n_row, _, r_h_r = h_core.shape
            r_psi_l, _, r_psi_r = psi_core.shape

            # Contract over n_col (physical index)
            # result[i,a, n_row, j,b] = sum_n H[i, n_row, n, j] * psi[a, n, b]
            contracted = torch.einsum('ionj,anb->iaojb', h_core, psi_core)
            # contracted has shape (r_h_l, r_psi_l, n_row, r_h_r, r_psi_r)

            # Reshape to (r_h_l * r_psi_l, n_row, r_h_r * r_psi_r)
            # But need to handle boundary conditions for ttm ranks
            if ttm_idx == 0:
                # First ttm core: r_h_l = 1, so result left rank = r_psi_l
                contracted = contracted.squeeze(0)  # (a, n_row, j, b)
                result.cores[k] = contracted.reshape(r_psi_l, n_row, r_h_r * r_psi_r)
            elif ttm_idx == ttm.d - 1:
                # Last ttm core: r_h_r = 1, so result right rank = r_psi_r
                contracted = contracted.squeeze(3)  # (i, a, n_row, b)
                result.cores[k] = contracted.reshape(r_h_l * r_psi_l, n_row, r_psi_r)
            else:
                # Middle core: both ranks are products
                result.cores[k] = contracted.reshape(r_h_l * r_psi_l, n_row, r_h_r * r_psi_r)
        else:
            # After active region: copy tt core
            result.cores[k] = tt.cores[k].clone()

    return result


def _apply_general(
    ttm: TTMatrix,
    tt: TensorTrain,
    dims_ttm: List[int],
    dims_tt: List[int],
) -> TensorTrain:
    """
    Apply TTMatrix to arbitrary (possibly non-contiguous) dimensions.

    This handles the general case by tracking TTMatrix rank flow through
    passive (non-active) dimensions.
    """
    # Map tt dimension -> ttm index (or None if not active)
    tt_dim_to_ttm = {}
    for ttm_idx, tt_dim in zip(dims_ttm, dims_tt):
        tt_dim_to_ttm[tt_dim] = ttm_idx

    # Sort active dims to process in order
    active_dims_sorted = sorted(dims_tt)

    # Build result dimensions
    result_dims = list(tt.dims)
    for ttm_idx, tt_dim in zip(dims_ttm, dims_tt):
        result_dims[tt_dim] = ttm.row_dims[ttm_idx]

    # For non-contiguous case, we need to track which ttm rank is "in flight"
    # When we encounter an active dim, we contract with ttm
    # When we encounter a passive dim between active dims, we pass ttm rank through

    # Find ranges: before first active, between actives, after last active
    first_active = active_dims_sorted[0]
    last_active = active_dims_sorted[-1]

    # Track which ttm cores we've used (in order of dims_tt sorted)

    # We need to know which ttm core corresponds to each active dim
    # But ttm cores are ordered by dims_ttm, not by dims_tt sorted order
    # For now, assume ttm dims are [0, 1, ..., d_m-1] and dims_tt gives the mapping
    # active_dims_sorted[i] in tt corresponds to dims_ttm[dims_tt.index(active_dims_sorted[i])]

    def get_ttm_idx_for_tt_dim(tt_dim):
        """Get the ttm core index for a given tt dimension."""
        pos = dims_tt.index(tt_dim)
        return dims_ttm[pos]

    # Compute result ranks
    result_ranks = []
    for k in range(tt.d - 1):
        if k < first_active:
            # Before any active dim
            result_ranks.append(tt.ranks[k + 1])
        elif k >= last_active:
            # After all active dims
            result_ranks.append(tt.ranks[k + 1])
        else:
            # Between active dims (or at an active dim boundary)
            # Need to track ttm rank
            # Find the most recent active dim <= k
            recent_active = max(d for d in active_dims_sorted if d <= k)
            ttm_idx = get_ttm_idx_for_tt_dim(recent_active)
            ttm_right_rank = ttm.ranks[ttm_idx + 1]

            if ttm_right_rank == 1:
                # TTM has boundary rank 1, no extra rank to carry
                result_ranks.append(tt.ranks[k + 1])
            else:
                # Need to carry ttm rank through
                result_ranks.append(tt.ranks[k + 1] * ttm_right_rank)

    result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                         tt.dtype, tt.device)

    # Now build cores
    for k in range(tt.d):
        psi_core = tt.cores[k]  # (r_psi_l, n, r_psi_r)
        r_psi_l, n_psi, r_psi_r = psi_core.shape

        if k in tt_dim_to_ttm:
            # Active dimension: contract with ttm
            ttm_idx = tt_dim_to_ttm[k]
            h_core = ttm.cores[ttm_idx]  # (r_h_l, n_row, n_col, r_h_r)
            r_h_l, n_row, _, r_h_r = h_core.shape

            # Contract over n_col (physical index)
            # result[i,a, n_row, j,b] = sum_n H[i, n_row, n, j] * psi[a, n, b]
            contracted = torch.einsum('ionj,anb->iaojb', h_core, psi_core)
            # contracted has shape (r_h_l, r_psi_l, n_row, r_h_r, r_psi_r)

            # Determine output shape based on position
            if k == first_active and k == last_active:
                # Only active dim: squeeze both ttm ranks
                contracted = contracted.squeeze(0).squeeze(2)  # (a, n_row, b)
                result.cores[k] = contracted.reshape(r_psi_l, n_row, r_psi_r)
            elif k == first_active:
                # First active: squeeze left ttm rank
                contracted = contracted.squeeze(0)  # (a, n_row, j, b)
                result.cores[k] = contracted.reshape(r_psi_l, n_row, r_h_r * r_psi_r)
            elif k == last_active:
                # Last active: squeeze right ttm rank
                contracted = contracted.squeeze(3)  # (i, a, n_row, b)
                result.cores[k] = contracted.reshape(r_h_l * r_psi_l, n_row, r_psi_r)
            else:
                # Middle active
                result.cores[k] = contracted.reshape(r_h_l * r_psi_l, n_row, r_h_r * r_psi_r)

        elif k < first_active or k > last_active:
            # Outside active region: just copy
            result.cores[k] = psi_core.clone()
        else:
            # Passive dim between active dims: need to pass ttm rank through
            # Create identity tensor to carry ttm rank
            # result[i,a, n, i,b] = delta_{ii'} * psi[a, n, b] (but we just copy i through)

            # Find the ttm rank we're carrying
            recent_active = max(d for d in active_dims_sorted if d < k)
            ttm_idx = get_ttm_idx_for_tt_dim(recent_active)
            r_h_carry = ttm.ranks[ttm_idx + 1]

            if r_h_carry == 1:
                # No rank to carry, just copy
                result.cores[k] = psi_core.clone()
            else:
                # Tensor product with identity: psi ⊗ I_{r_h}
                # psi: (r_psi_l, n, r_psi_r)
                # I: (r_h, r_h)
                # result: (r_psi_l * r_h, n, r_psi_r * r_h)
                I_h = torch.eye(r_h_carry, dtype=tt.dtype, device=tt.device)
                # result[a,i, n, b,j] = psi[a,n,b] * I[i,j]
                result_core = torch.einsum('anb,ij->ainbj', psi_core, I_h)
                result.cores[k] = result_core.reshape(r_psi_l * r_h_carry, n_psi, r_psi_r * r_h_carry)

    return result


def matmul(ttm1: TTMatrix, ttm2: TTMatrix) -> TTMatrix:
    """
    Compute the matrix product of two TTMatrices.

    Computes result = ttm1 @ ttm2, where:
    - ttm1 maps col_dims_1 -> row_dims_1
    - ttm2 maps col_dims_2 -> row_dims_2
    - result maps col_dims_2 -> row_dims_1

    Requires ttm1.col_dims == ttm2.row_dims (compatible for multiplication).

    Result ranks are products of input ranks.

    Args:
        ttm1: Left TTMatrix operator
        ttm2: Right TTMatrix operator

    Returns:
        TTMatrix representing ttm1 @ ttm2

    Example:
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])
        H12 = matmul(H1, H2)  # ranks become [1, 6, 1]
    """
    # Validate dimensions
    if ttm1.d != ttm2.d:
        raise ValueError(
            f"TTMatrices must have same number of sites: {ttm1.d} vs {ttm2.d}"
        )

    if ttm1.col_dims != ttm2.row_dims:
        raise ValueError(
            f"Incompatible dimensions for matmul: "
            f"ttm1.col_dims={ttm1.col_dims} != ttm2.row_dims={ttm2.row_dims}"
        )

    # Result dimensions: row_dims from ttm1, col_dims from ttm2
    result_dims = [(ttm1.row_dims[k], ttm2.col_dims[k]) for k in range(ttm1.d)]

    # Result ranks are products of input ranks
    result_ranks = [ttm1.ranks[k] * ttm2.ranks[k] for k in range(1, ttm1.d)]

    result = TTMatrix(result_dims, result_ranks if result_ranks else None,
                      ttm1.dtype, ttm1.device)

    # Contract cores at each site
    for k in range(ttm1.d):
        core1 = ttm1.cores[k]  # (r1_l, n_row, n_mid, r1_r)
        core2 = ttm2.cores[k]  # (r2_l, n_mid, n_col, r2_r)

        r1_l, n_row, _, r1_r = core1.shape
        r2_l, _, n_col, r2_r = core2.shape

        # Contract over n_mid (the shared dimension)
        # result[i,a, n_row, n_col, j,b] = sum_n_mid core1[i, n_row, n_mid, j] * core2[a, n_mid, n_col, b]
        contracted = torch.einsum('imnj,anob->iamojb', core1, core2)

        # Reshape to (r1_l * r2_l, n_row, n_col, r1_r * r2_r)
        result.cores[k] = contracted.reshape(r1_l * r2_l, n_row, n_col, r1_r * r2_r)

    return result


def transpose(ttm: TTMatrix) -> TTMatrix:
    """
    Compute the transpose of a TTMatrix.

    Swaps row and column dimensions at each site:
    - Original: maps col_dims -> row_dims
    - Transposed: maps row_dims -> col_dims

    Args:
        ttm: TTMatrix to transpose

    Returns:
        TTMatrix representing ttm^T

    Example:
        H = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])  # (15, 24)
        Ht = transpose(H)  # (24, 15)
    """
    # Swap row and column dims
    transposed_dims = [(n_col, n_row) for n_row, n_col in ttm.dims]

    result = TTMatrix(transposed_dims, ttm.ranks[1:-1], ttm.dtype, ttm.device)

    # Transpose each core: (r_l, n_row, n_col, r_r) -> (r_l, n_col, n_row, r_r)
    for k in range(ttm.d):
        # Swap axes 1 and 2 (row and col)
        result.cores[k] = ttm.cores[k].permute(0, 2, 1, 3).contiguous()

    return result


def adjoint(ttm: TTMatrix) -> TTMatrix:
    """
    Compute the adjoint (conjugate transpose) of a TTMatrix.

    For real matrices, this is the same as transpose.
    For complex matrices, this is the conjugate transpose (Hermitian adjoint).

    Args:
        ttm: TTMatrix to adjoint

    Returns:
        TTMatrix representing ttm^H (or ttm^T for real)

    Example:
        H = TTMatrix.random([(3, 4), (5, 6)], ranks=[2])
        Hdag = adjoint(H)  # H^H
    """
    # Swap row and column dims
    adjoint_dims = [(n_col, n_row) for n_row, n_col in ttm.dims]

    result = TTMatrix(adjoint_dims, ttm.ranks[1:-1], ttm.dtype, ttm.device)

    # Transpose and conjugate each core: (r_l, n_row, n_col, r_r) -> (r_l, n_col, n_row, r_r)
    for k in range(ttm.d):
        # Swap axes 1 and 2 (row and col), then conjugate
        transposed = ttm.cores[k].permute(0, 2, 1, 3).contiguous()
        if transposed.is_complex():
            result.cores[k] = transposed.conj()
        else:
            result.cores[k] = transposed

    return result


def trace(ttm: TTMatrix) -> float:
    """
    Compute the trace of a square TTMatrix.

    The trace is the sum of diagonal elements. Requires n_row == n_col at each site.

    Args:
        ttm: Square TTMatrix (n_row == n_col at each site)

    Returns:
        Scalar trace value

    Raises:
        ValueError: If TTMatrix is not square at all sites

    Example:
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        tr = trace(H)  # sum of diagonal of 12x12 matrix
    """
    # Validate square at each site
    for k, (n_row, n_col) in enumerate(ttm.dims):
        if n_row != n_col:
            raise ValueError(
                f"TTMatrix must be square at each site for trace. "
                f"Site {k} has dims ({n_row}, {n_col})"
            )

    # Contract diagonal at each site: sum over i where row_idx == col_idx
    # Start with boundary matrix of shape (1, 1)
    result = torch.ones(1, 1, dtype=ttm.dtype, device=ttm.device)

    for k in range(ttm.d):
        core = ttm.cores[k]  # (r_l, n, n, r_r)

        # Contract diagonal: M[i, j] = sum_a core[i, a, a, j]
        # Use einsum to sum over diagonal of physical indices
        M = torch.einsum('iaaj->ij', core)  # (r_l, r_r)

        # Chain multiply
        result = result @ M

    return result.item()


def kron(*ttms: TTMatrix) -> TTMatrix:
    """
    Kronecker (tensor) product of TTMatrices.

    Accepts two or more TTMatrices. The result acts on the concatenated sites,
    with the first operand's sites first, then the second's, etc.

    In matrix terms: if H1 is (m1, n1) and H2 is (m2, n2), then
    kron(H1, H2) is (m1*m2, n1*n2).

    Args:
        *ttms: Two or more TTMatrices

    Returns:
        TTMatrix representing ttm1 ⊗ ttm2 ⊗ ... ⊗ ttmN

    Example:
        # Build a multi-site operator from single-site operators
        X = TTMatrix.local(pauli_x)
        Z = TTMatrix.local(pauli_z)

        XZ = kron(X, Z)           # X ⊗ Z, acts on 2 sites
        XZXZ = kron(X, Z, X, Z)   # 4-site operator

        # Or use & operator
        XZXZ = X & Z & X & Z

        # Apply to specific sites
        result = XZXZ.on(0, 2, 4, 6) @ psi
    """
    if len(ttms) < 2:
        raise ValueError("kron requires at least 2 TTMatrices")

    # Binary kron implementation
    def _kron_binary(ttm1: TTMatrix, ttm2: TTMatrix) -> TTMatrix:
        # Result dimensions: concatenate dims
        result_dims = list(ttm1.dims) + list(ttm2.dims)

        # Result internal ranks: ttm1's internal ranks, then junction rank 1, then ttm2's internal
        result_internal_ranks = list(ttm1.ranks[1:-1]) + [1] + list(ttm2.ranks[1:-1])

        # Handle edge case where one or both have d=1 (no internal ranks)
        if ttm1.d == 1 and ttm2.d == 1:
            result_internal_ranks = [1]
        elif ttm1.d == 1:
            result_internal_ranks = [1] + list(ttm2.ranks[1:-1])
        elif ttm2.d == 1:
            result_internal_ranks = list(ttm1.ranks[1:-1]) + [1]

        result = TTMatrix(result_dims, result_internal_ranks, ttm1.dtype, ttm1.device)

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

        # Copy cores from ttm2
        for k in range(ttm2.d):
            result.cores[ttm1.d + k] = ttm2.cores[k].clone()

        return result

    # Fold left: kron(A, B, C, D) = kron(kron(kron(A, B), C), D)
    result = ttms[0]
    for ttm in ttms[1:]:
        result = _kron_binary(result, ttm)

    return result


# ============================================================================
# TTMatrix operations via TensorTrain conversion
# ============================================================================
# These operations work on the bond structure and are agnostic to the
# physical index interpretation. We convert to TensorTrain, use the
# existing implementation, and convert back.


def add(ttm1: TTMatrix, ttm2: TTMatrix) -> TTMatrix:
    """
    Add two TTMatrices.

    Result ranks are sums of input ranks (before truncation).

    Args:
        ttm1: First TTMatrix
        ttm2: Second TTMatrix (must have same dims)

    Returns:
        TTMatrix representing ttm1 + ttm2

    Example:
        H1 = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])
        H_sum = add(H1, H2)  # ranks become [1, 5, 1]
    """
    if ttm1.dims != ttm2.dims:
        raise ValueError(
            f"TTMatrix dims must match: {ttm1.dims} vs {ttm2.dims}"
        )

    from .tensortrain import add as tt_add

    tt1 = ttm1.to_tensor_train()
    tt2 = ttm2.to_tensor_train()
    tt_result = tt_add(tt1, tt2)

    return TTMatrix.from_tensor_train(tt_result, ttm1.dims)


def scale(ttm: TTMatrix, alpha: float) -> TTMatrix:
    """
    Scale a TTMatrix by a scalar.

    Args:
        ttm: TTMatrix to scale
        alpha: Scalar multiplier

    Returns:
        TTMatrix representing alpha * ttm

    Example:
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        H2 = scale(H, 2.0)  # 2 * H
    """
    from .tensortrain import scale as tt_scale

    tt = ttm.to_tensor_train()
    tt_result = tt_scale(tt, alpha)

    return TTMatrix.from_tensor_train(tt_result, ttm.dims)


def norm_squared(ttm: TTMatrix) -> float:
    """
    Compute the squared Frobenius norm of a TTMatrix.

    ||A||_F^2 = sum_{i,j} |A_{i,j}|^2 = tr(A^H A)

    Args:
        ttm: TTMatrix

    Returns:
        Squared Frobenius norm as scalar

    Example:
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        ns = norm_squared(H)
    """
    from . import linalg

    tt = ttm.to_tensor_train()
    return linalg.norm_squared(tt)


def norm(ttm: TTMatrix) -> float:
    """
    Compute the Frobenius norm of a TTMatrix.

    ||A||_F = sqrt(sum_{i,j} |A_{i,j}|^2)

    Args:
        ttm: TTMatrix

    Returns:
        Frobenius norm as scalar

    Example:
        H = TTMatrix.random([(3, 3), (4, 4)], ranks=[2])
        n = norm(H)
    """
    return norm_squared(ttm) ** 0.5


def round(
    ttm: TTMatrix,
    max_rank: Optional[int] = None,
    rel_tol: Optional[float] = None,
    abs_tol: Optional[float] = None,
) -> TTMatrix:
    """
    Round (truncate) a TTMatrix to lower rank via SVD.

    Args:
        ttm: TTMatrix to truncate
        max_rank: Maximum rank to keep
        rel_tol: Relative tolerance for singular values
        abs_tol: Absolute tolerance for singular values

    Returns:
        TTMatrix with reduced ranks

    Example:
        H = add(H1, H2)  # ranks grew
        H_rounded = round(H, max_rank=5)  # truncate to rank 5
    """
    from . import linalg

    tt = ttm.to_tensor_train()
    tt_result = linalg.svd(tt, max_rank=max_rank, rel_tol=rel_tol, abs_tol=abs_tol)

    return TTMatrix.from_tensor_train(tt_result, ttm.dims)


def kronadd(
    ttm1: TTMatrix,
    ttm2: TTMatrix,
    sites: Tuple[List[int], List[int]],
) -> TTMatrix:
    """
    Add two TTMatrices acting on different but overlapping sites.

    Computes (ttm1 ⊗ I) + (I ⊗ ttm2) where identities fill non-active sites.
    This is the natural operation for building Hamiltonians from local terms.

    Note: Unlike TensorTrain.kronadd which uses ones-tensors for inactive paths,
    TTMatrix.kronadd uses identity matrices, which is the correct semantic for
    operators (Hamiltonians, etc.).

    Args:
        ttm1: First TTMatrix with d1 sites
        ttm2: Second TTMatrix with d2 sites
        sites: Tuple of two lists specifying which result sites each operand acts on.
               Both lists must be contiguous, and their union must be contiguous (no gaps).

    Returns:
        TTMatrix acting on the union of sites

    Constraints:
        - sites[0] must be contiguous: [a, a+1, ..., a+d1-1]
        - sites[1] must be contiguous: [b, b+1, ..., b+d2-1]
        - Union must be contiguous (no gap): max(a, b) <= min(a+d1, b+d2)
        - At overlapping sites, dimensions must match

    Example:
        # Build H = H_01 + H_12 acting on sites [0, 1, 2]
        H_01 = TTMatrix.random([(2, 2), (3, 3)], ranks=[2])
        H_12 = TTMatrix.random([(3, 3), (4, 4)], ranks=[3])
        H = kronadd(H_01, H_12, sites=([0, 1], [1, 2]))

        # Non-overlapping but adjacent: H = H_01 + H_23
        H = kronadd(H_01, H_23, sites=([0, 1], [2, 3]))

        # Invalid - gap between sites:
        # kronadd(H_01, H_34, sites=([0, 1], [3, 4]))  # Error!

        # Invalid - non-contiguous sites:
        # kronadd(H, H, sites=([0, 2], [1, 3]))  # Error!
    """
    sites1, sites2 = list(sites[0]), list(sites[1])

    # Validate site lengths match operators
    if len(sites1) != ttm1.d:
        raise ValueError(
            f"sites[0] length ({len(sites1)}) must match ttm1 sites ({ttm1.d})"
        )
    if len(sites2) != ttm2.d:
        raise ValueError(
            f"sites[1] length ({len(sites2)}) must match ttm2 sites ({ttm2.d})"
        )

    # Check contiguity of sites1
    for i in range(1, len(sites1)):
        if sites1[i] != sites1[i - 1] + 1:
            raise ValueError(
                f"sites[0] must be contiguous, got {sites1}"
            )

    # Check contiguity of sites2
    for i in range(1, len(sites2)):
        if sites2[i] != sites2[i - 1] + 1:
            raise ValueError(
                f"sites[1] must be contiguous, got {sites2}"
            )

    # Compute ranges
    a_start, a_end = sites1[0], sites1[-1]  # inclusive
    b_start, b_end = sites2[0], sites2[-1]  # inclusive

    # Check no gap between ranges (union must be contiguous)
    if min(a_end, b_end) < max(a_start, b_start) - 1:
        raise ValueError(
            f"sites must have contiguous union (no gaps). "
            f"sites[0]=[{a_start}..{a_end}], sites[1]=[{b_start}..{b_end}]"
        )

    # Result site range
    result_start = min(a_start, b_start)
    result_end = max(a_end, b_end)
    result_d = result_end - result_start + 1

    # Build result dimensions and validate overlaps
    result_dims = []
    for k in range(result_d):
        result_site = result_start + k
        in_a = a_start <= result_site <= a_end
        in_b = b_start <= result_site <= b_end

        if in_a and in_b:
            # Overlap: dimensions must match
            a_idx = result_site - a_start
            b_idx = result_site - b_start
            if ttm1.dims[a_idx] != ttm2.dims[b_idx]:
                raise ValueError(
                    f"Dimension mismatch at overlapping site {result_site}: "
                    f"ttm1.dims[{a_idx}]={ttm1.dims[a_idx]} != "
                    f"ttm2.dims[{b_idx}]={ttm2.dims[b_idx]}"
                )
            result_dims.append(ttm1.dims[a_idx])
        elif in_a:
            a_idx = result_site - a_start
            result_dims.append(ttm1.dims[a_idx])
        else:  # in_b
            b_idx = result_site - b_start
            result_dims.append(ttm2.dims[b_idx])

    # Compute result ranks using block structure
    # At each bond, both operators need a "path" - either their internal rank
    # if active, or rank 1 for the identity path if waiting/terminated.
    result_ranks = []
    for k in range(result_d - 1):
        result_site = result_start + k
        # Bond is between result_site and result_site + 1

        # ttm1's contribution to this bond
        if result_site < a_start:
            r_a = 1  # Before ttm1 starts: identity path
        elif result_site >= a_end:
            r_a = 1  # ttm1 has terminated: identity path
        else:
            a_idx = result_site - a_start
            r_a = ttm1.ranks[a_idx + 1]

        # ttm2's contribution to this bond
        if result_site < b_start:
            r_b = 1  # Before ttm2 starts: identity path
        elif result_site >= b_end:
            r_b = 1  # ttm2 has terminated: identity path
        else:
            b_idx = result_site - b_start
            r_b = ttm2.ranks[b_idx + 1]

        result_ranks.append(r_a + r_b)

    # Create result TTMatrix
    result = TTMatrix(result_dims, result_ranks if result_ranks else None,
                      ttm1.dtype, ttm1.device)

    # Determine which operator starts first (comes first in block structure)
    a_first = a_start <= b_start

    # Build cores using block structure
    for k in range(result_d):
        result_site = result_start + k
        in_a = a_start <= result_site <= a_end
        in_b = b_start <= result_site <= b_end

        n_row, n_col = result_dims[k]

        # Get actual left and right ranks for this core
        r_left = 1 if k == 0 else result_ranks[k - 1]
        r_right = 1 if k == result_d - 1 else result_ranks[k]

        # Helper function to get operator's rank contribution at a bond
        def get_rank_contribution(op_start, op_end, op_ranks, site, is_left):
            if is_left:
                if k == 0:
                    return 0
                prev = site - 1
                if prev < op_start or prev >= op_end:
                    return 1  # Identity path
                else:
                    idx = prev - op_start
                    return op_ranks[idx + 1]
            else:
                if k == result_d - 1:
                    return 0
                if site < op_start or site >= op_end:
                    return 1  # Identity path
                else:
                    idx = site - op_start
                    return op_ranks[idx + 1]

        # Get rank contributions for block positioning
        if a_first:
            r_first_left = get_rank_contribution(a_start, a_end, ttm1.ranks, result_site, True)
            r_first_right = get_rank_contribution(a_start, a_end, ttm1.ranks, result_site, False)
        else:
            r_first_left = get_rank_contribution(b_start, b_end, ttm2.ranks, result_site, True)
            r_first_right = get_rank_contribution(b_start, b_end, ttm2.ranks, result_site, False)

        # Initialize core to zeros
        core = torch.zeros(r_left, n_row, n_col, r_right, dtype=ttm1.dtype, device=ttm1.device)

        # Helper to place operator block
        def place_operator_block(op_core, is_first):
            rL, _, _, rR = op_core.shape

            if is_first:
                left_offset = 0
                right_offset = 0
            else:
                left_offset = r_first_left
                right_offset = r_first_right

            # At first site, all share left index 0
            if k == 0:
                left_offset = 0
                if not is_first and (in_a and in_b):
                    right_offset = r_first_right

            # At last site, all share right index 0
            if k == result_d - 1:
                right_offset = 0

            core[left_offset:left_offset + rL, :, :,
                 right_offset:right_offset + rR] = op_core

        # Place ttm1's block if active
        if in_a:
            a_idx = result_site - a_start
            a_core = ttm1.cores[a_idx]
            place_operator_block(a_core, a_first)

        # Place ttm2's block if active
        if in_b:
            b_idx = result_site - b_start
            b_core = ttm2.cores[b_idx]
            place_operator_block(b_core, not a_first)

        # Handle identity blocks for paths that are "waiting" or "done"
        # Identity matrix in (1, n_row, n_col, 1) shape
        identity = torch.eye(n_row, n_col, dtype=ttm1.dtype, device=ttm1.device)
        identity_reshaped = identity.reshape(1, n_row, n_col, 1)

        # Determine first/second based on which starts first
        if a_first:
            first_end, second_start = a_end, b_start
            first_active, second_active = in_a, in_b
        else:
            first_end, second_start = b_end, a_start
            first_active, second_active = in_b, in_a

        # Case: First operator has terminated but we're in second operator's region
        if not first_active and second_active and result_site > first_end:
            core[0:1, :, :, 0:1] = identity_reshaped

        # Case: Second operator hasn't started but first is active
        if not second_active and first_active and result_site < second_start:
            if k == 0:
                right_pos = r_first_right
                core[0:1, :, :, right_pos:right_pos + 1] = identity_reshaped
            elif k < result_d - 1:
                left_pos = r_first_left
                right_pos = r_first_right
                core[left_pos:left_pos + 1, :, :, right_pos:right_pos + 1] = identity_reshaped

        result.cores[k] = core

    return result
