"""Linear algebra operations for tensor trains."""

import torch
from typing import Tuple, Optional, List, Union
from .tensortrain import TensorTrain, dot, hadamard

# Re-export TTMatrix operations for backward compatibility
from .ttmatrix import (
    apply, matmul, transpose, adjoint, trace, kron,
    add as ttm_add, scale as ttm_scale,
    norm as ttm_norm, norm_squared as ttm_norm_squared,
    round as ttm_round,
)


# ============================================================================
# Sweep Utilities
# ============================================================================


def _transpose_bonds(core: torch.Tensor) -> torch.Tensor:
    """
    Transpose bond dimensions of a TT core: (r_left, n, r_right) -> (r_right, n, r_left).

    This swaps the left and right bond dimensions while keeping the physical dimension.
    """
    return core.permute(2, 1, 0)


def _get_prepared_cores(
    tt: TensorTrain,
    k: int,
    direction: str
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    """
    Get and prepare two neighboring cores to look like left-to-right configuration.

    This is a UNIVERSAL helper for any operation on two neighboring cores.
    Maps the uniform index k (always 0 to d-1) to actual positions based on direction.
    After preparation, algorithms can always assume left-to-right layout.

    Args:
        tt: Tensor train
        k: Sweep index (always 0 to d-1)
        direction: 'left-to-right' or 'right-to-left'

    Returns:
        Tuple of (prepared_current, prepared_neighbor)
        - For left-to-right: returns cores as-is from positions k and k+1
        - For right-to-left: returns transposed cores from positions d-1-k and d-2-k

    Raises:
        ValueError: If direction is not 'left-to-right' or 'right-to-left'
    """
    if direction not in ('left-to-right', 'right-to-left'):
        raise ValueError(f"Invalid direction: {direction}. Use 'left-to-right' or 'right-to-left'")

    if direction == 'left-to-right':
        position = k
        neighbor_pos = k + 1 if k < tt.d - 1 else None
    else:  # right-to-left
        position = tt.d - 1 - k
        neighbor_pos = tt.d - 2 - k if k < tt.d - 1 else None

    current = tt.cores[position]
    neighbor = tt.cores[neighbor_pos] if neighbor_pos is not None else None

    if direction == 'left-to-right':
        return current, neighbor
    else:
        # Transpose bonds to make right-to-left look like left-to-right
        prep_current = _transpose_bonds(current)
        prep_neighbor = _transpose_bonds(neighbor) if neighbor is not None else None
        return prep_current, prep_neighbor


def _unprepare_and_update_cores(
    tt: TensorTrain,
    k: int,
    current: torch.Tensor,
    neighbor: Optional[torch.Tensor],
    direction: str
) -> None:
    """
    Reverse the preparation transformation and update cores in the tensor train.

    This is a UNIVERSAL helper that reverses _get_prepared_cores and writes back.
    Maps the uniform index k to actual positions and updates both cores and ranks.

    Args:
        tt: Tensor train to update
        k: Sweep index (always 0 to d-1)
        current: Processed current core
        neighbor: Processed neighbor core (or None)
        direction: 'left-to-right' or 'right-to-left'
    """
    if direction == 'left-to-right':
        position = k
        neighbor_pos = k + 1 if k < tt.d - 1 else None

        # Cores are already in correct orientation
        tt.cores[position] = current
        if neighbor_pos is not None:
            tt.cores[neighbor_pos] = neighbor
            # Update rank: right rank of current core
            tt.ranks[neighbor_pos] = current.shape[2]
        else:
            # Last core - update final rank
            tt.ranks[tt.d] = current.shape[2]
    else:  # right-to-left
        position = tt.d - 1 - k
        neighbor_pos = tt.d - 2 - k if k < tt.d - 1 else None

        # Transpose back to original orientation
        final_current = _transpose_bonds(current)
        final_neighbor = _transpose_bonds(neighbor) if neighbor is not None else None

        tt.cores[position] = final_current
        if neighbor_pos is not None:
            tt.cores[neighbor_pos] = final_neighbor
            # Update rank: left rank of current core
            tt.ranks[position] = final_current.shape[0]
        else:
            # Last core - update first rank
            tt.ranks[0] = final_current.shape[0]


# ============================================================================
# Local Decomposition Operations
# ============================================================================


def qr_step(core: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Perform local QR decomposition on a single TT core (left-to-right convention).

    This function ONLY knows about left-to-right decomposition. Direction handling
    is done by prepare/unprepare helpers that transpose cores as needed.

    Args:
        core: Tensor of shape (r_left, n, r_right) in left-to-right convention

    Returns:
        Tuple of (Q_core, R) where:
            - Q_core: Left-orthogonal core of shape (r_left, n, r_new)
            - R: Upper triangular matrix of shape (r_new, r_right) to contract into right neighbor
    """
    r_left, n, r_right = core.shape

    # Matricize as (r_left * n, r_right) for left-orthogonal decomposition
    core_mat = core.reshape(r_left * n, r_right)

    # Perform QR decomposition
    Q, R = torch.linalg.qr(core_mat)

    # Reshape Q back to core format
    r_new = Q.shape[1]
    Q_core = Q.reshape(r_left, n, r_new)

    return Q_core, R


def qr(tt: TensorTrain, direction: str = 'left-to-right') -> Tuple[TensorTrain, torch.Tensor]:
    """
    Perform QR/LQ decomposition of a tensor train with directional sweep.

    Direction-dependent logic is FULLY encapsulated in universal prepare/unprepare
    helpers. The main loop is completely uniform regardless of direction.

    Args:
        tt: Input tensor train
        direction: Sweep direction - 'left-to-right' or 'right-to-left'
            - 'left-to-right': Produces left-orthogonal cores (Q^T Q = I)
            - 'right-to-left': Produces right-orthogonal cores (Q Q^T = I)

    Returns:
        Tuple of (qr_tt, residual) where:
            - qr_tt: TensorTrain with orthogonal cores
            - residual: Final triangular matrix from the last core in sweep
    """
    qr_tt = tt.copy()

    for k in range(qr_tt.d):
        # Get and prepare cores to look like left-to-right
        prep_current, prep_neighbor = _get_prepared_cores(qr_tt, k, direction)

        # QR decomposition (always left-to-right after preparation)
        Q_core, R = qr_step(prep_current)

        # Contract R into neighbor if it exists
        if prep_neighbor is not None:
            # R @ neighbor: (r_new, r_old) @ (r_old, n, r_next) -> (r_new, n, r_next)
            prep_neighbor = torch.einsum('ij,jkl->ikl', R, prep_neighbor)

        # Unprepare and update cores
        _unprepare_and_update_cores(qr_tt, k, Q_core, prep_neighbor, direction)

    return qr_tt, R


def svd_step(
    core: torch.Tensor,
    max_rank: Optional[int] = None,
    rel_tol: Optional[float] = None,
    abs_tol: Optional[float] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Perform local SVD decomposition on a single TT core (left-to-right convention).

    This function ONLY knows about left-to-right decomposition. Direction handling
    is done by prepare/unprepare helpers that transpose cores as needed.

    Args:
        core: Tensor of shape (r_left, n, r_right) in left-to-right convention
        max_rank: Maximum rank to keep (optional)
        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:
        Tuple of (U_core, S, SVt) where:
            - U_core: Left-orthogonal core (r_left, n, r_new)
            - S: Singular values (r_new,)
            - SVt: S @ V^T residual of shape (r_new, r_right) to contract into right neighbor
    """
    r_left, n, r_right = core.shape

    # Matricize as (r_left * n, r_right) for left-orthogonal decomposition
    core_mat = core.reshape(r_left * n, r_right)

    # Perform SVD
    U, S, Vt = torch.linalg.svd(core_mat, full_matrices=False)

    # Determine rank to keep
    r_new = len(S)

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

    # Apply relative tolerance truncation
    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()))

    # Apply absolute tolerance truncation
    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, :]

    # U is the orthogonal core, S @ Vt is the residual
    U_core = U.reshape(r_left, n, r_new)
    SVt = S.unsqueeze(1) * Vt  # (r_new, r_right)

    return U_core, S, SVt


def svd(
    tt: TensorTrain,
    max_rank: Optional[int] = None,
    rel_tol: Optional[float] = None,
    abs_tol: Optional[float] = None,
    direction: str = 'left-to-right',
) -> TensorTrain:
    """
    Perform TT-SVD decomposition with rank truncation and directional sweep.

    Direction-dependent logic is FULLY encapsulated in universal prepare/unprepare
    helpers. The main loop is completely uniform regardless of direction.

    Applies SVD to each core sequentially, propagating residual factors in the
    sweep direction. The result is an orthogonalized tensor train with potentially
    reduced ranks.

    Args:
        tt: Input tensor train
        max_rank: Maximum rank to keep at each step (optional)
        rel_tol: Relative tolerance for singular value truncation (optional)
        abs_tol: Absolute tolerance for singular value truncation (optional)
        direction: Sweep direction - 'left-to-right' or 'right-to-left'
            - 'left-to-right': Produces left-orthogonal cores
            - 'right-to-left': Produces right-orthogonal cores

    Returns:
        TensorTrain with orthogonal cores and potentially reduced ranks
    """
    svd_tt = tt.copy()

    for k in range(svd_tt.d):
        # Get and prepare cores to look like left-to-right
        prep_current, prep_neighbor = _get_prepared_cores(svd_tt, k, direction)

        # SVD decomposition (always left-to-right after preparation)
        U_core, _, SVt = svd_step(prep_current, max_rank, rel_tol, abs_tol)

        # Contract SVt into neighbor if it exists
        if prep_neighbor is not None:
            # SVt @ neighbor: (r_new, r_old) @ (r_old, n, r_next) -> (r_new, n, r_next)
            prep_neighbor = torch.einsum('ij,jkl->ikl', SVt, prep_neighbor)
        else:
            # Last core in sweep - absorb residual back into core to preserve tensor
            r_left, n, _ = U_core.shape
            U_mat = U_core.reshape(r_left * n, -1)
            # Reconstruct: U @ diag(S) @ Vt where SVt = diag(S) @ Vt
            prep_current = (U_mat @ SVt).reshape(r_left, n, -1)
            U_core = prep_current

        # Unprepare and update cores
        _unprepare_and_update_cores(svd_tt, k, U_core, prep_neighbor, direction)

    return svd_tt


def norm_squared(tt: TensorTrain) -> float:
    """
    Compute the squared Frobenius norm of a tensor train.

    Uses the efficient TT dot product to compute ||tt||² = <tt, tt>
    without forming the full tensor. Complexity is O(n·r³) instead of
    O(n^d) for the full tensor.

    Args:
        tt: Input tensor train

    Returns:
        Squared Frobenius norm as a scalar float
    """
    # Compute <tt, tt> via boundary matrices
    matrices = dot(tt, tt)
    # The last boundary matrix is (1, 1) containing the inner product
    return matrices[-1].item()


def norm(tt: TensorTrain) -> float:
    """
    Compute the Frobenius norm of a tensor train.

    Uses the efficient TT dot product to compute ||tt|| = sqrt(<tt, tt>)
    without forming the full tensor. Complexity is O(n·r³) instead of
    O(n^d) for the full tensor.

    Args:
        tt: Input tensor train

    Returns:
        Frobenius norm as a scalar float
    """
    return norm_squared(tt) ** 0.5


# ============================================================================
# Reduction Operations
# ============================================================================


def sum(tt: TensorTrain, dim: Optional[int] = None):
    """
    Sum a tensor train over one or all dimensions.

    Args:
        tt: Input tensor train
        dim: Dimension to sum over. If None, sums over all dimensions
             and returns a scalar.

    Returns:
        If dim is None: scalar float (sum of all elements)
        If dim is specified: TensorTrain with that dimension removed
    """
    if dim is None:
        return _sum_all(tt)
    else:
        return _sum_dim(tt, dim)


def _sum_all(tt: TensorTrain) -> float:
    """
    Sum over all dimensions of a tensor train.

    Contracts each core with a ones vector over the physical index,
    then chain-multiplies the resulting matrices.

    Args:
        tt: Input tensor train

    Returns:
        Scalar sum of all tensor elements
    """
    # Start with identity matrix of size (1, 1)
    result = torch.ones(1, 1, dtype=tt.dtype, device=tt.device)

    for k in range(tt.d):
        core = tt.cores[k]  # (r_left, n, r_right)

        # Contract physical index with ones: sum over n
        # core[i, :, j].sum() -> M[i, j]
        M = core.sum(dim=1)  # (r_left, r_right)

        # Chain multiply: result @ M
        result = result @ M

    return result.item()


def _sum_dim(tt: TensorTrain, dim: int) -> TensorTrain:
    """
    Sum over a single dimension of a tensor train.

    Contracts the specified core with a ones vector, then merges the
    resulting matrix into a neighboring core.

    Args:
        tt: Input tensor train
        dim: Dimension to sum over (0 <= dim < tt.d)

    Returns:
        TensorTrain with dimension `dim` removed
    """
    if dim < 0 or dim >= tt.d:
        raise ValueError(f"dim={dim} out of range for TT with {tt.d} dimensions")

    if tt.d == 1:
        raise ValueError("Cannot reduce a 1-dimensional TT to 0 dimensions")

    # Contract core[dim] over physical index
    core = tt.cores[dim]  # (r_left, n, r_right)
    M = core.sum(dim=1)   # (r_left, r_right)

    # Build new TT with one fewer dimension
    new_dims = tt.dims[:dim] + tt.dims[dim + 1:]

    if dim == 0:
        # Merge M into the right neighbor (core[1])
        # M: (1, r_right), neighbor: (r_right, n, r_next)
        neighbor = tt.cores[1]
        # M @ neighbor over bond index: (1, r_right) @ (r_right, n, r_next) -> (1, n, r_next)
        new_first_core = torch.einsum('ij,jkl->ikl', M, neighbor)

        # Build new ranks (skip the contracted bond)
        new_ranks = [tt.ranks[k + 2] for k in range(tt.d - 2)]

        result = TensorTrain(new_dims, new_ranks if new_ranks else None, tt.dtype, tt.device)
        result.cores[0] = new_first_core
        for k in range(2, tt.d):
            result.cores[k - 1] = tt.cores[k].clone()

    elif dim == tt.d - 1:
        # Merge M into the left neighbor (core[d-2])
        # neighbor: (r_prev, n, r_left), M: (r_left, 1)
        neighbor = tt.cores[dim - 1]
        # neighbor @ M over bond index: (r_prev, n, r_left) @ (r_left, 1) -> (r_prev, n, 1)
        new_last_core = torch.einsum('ijk,kl->ijl', neighbor, M)

        # Build new ranks (skip the contracted bond)
        new_ranks = [tt.ranks[k + 1] for k in range(tt.d - 2)]

        result = TensorTrain(new_dims, new_ranks if new_ranks else None, tt.dtype, tt.device)
        for k in range(dim - 1):
            result.cores[k] = tt.cores[k].clone()
        result.cores[dim - 1] = new_last_core

    else:
        # Middle dimension: merge M into the right neighbor (core[dim+1])
        # M: (r_left, r_right), neighbor: (r_right, n, r_next)
        neighbor = tt.cores[dim + 1]
        merged_core = torch.einsum('ij,jkl->ikl', M, neighbor)

        # Build new ranks: remove rank at position dim+1
        new_ranks = []
        for k in range(tt.d):
            if k != dim and k != dim + 1:
                new_ranks.append(tt.ranks[k + 1] if k < dim else tt.ranks[k + 1])
        # Actually simpler: ranks are [1, r1, r2, ..., r_{d-1}, 1]
        # We remove dimension dim, so we skip ranks[dim] and ranks[dim+1] merge
        new_ranks = []
        for k in range(1, tt.d):
            if k == dim + 1:
                continue  # Skip this rank (it's absorbed)
            new_ranks.append(tt.ranks[k])

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

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

        # The merged core takes position dim
        result.cores[dim] = merged_core

        # Copy cores after dim+1
        for k in range(dim + 2, tt.d):
            result.cores[k - 1] = tt.cores[k].clone()

    return result


def contract(
    tt1: TensorTrain,
    tt2: TensorTrain,
    dims: List[int],
) -> Union[TensorTrain, float]:
    """
    Contract two tensor trains over specified dimensions.

    This is a generalization of dot and hadamard:
    - contract(tt1, tt2, dims=[]) is equivalent to hadamard(tt1, tt2)
    - contract(tt1, tt2, dims=[0,1,...,d-1]) is equivalent to dot(tt1, tt2)

    For density matrix computations, contract over the "environment" dimensions
    to obtain the reduced density matrix for the remaining "system" dimensions.

    Args:
        tt1: First tensor train
        tt2: Second tensor train (must have same shape as tt1)
        dims: List of dimensions to contract over (sum out)

    Returns:
        If all dims contracted: scalar float (inner product)
        Otherwise: TensorTrain over the non-contracted dimensions

    Example:
        # For tensors of shape (4, 5, 6), contract over middle dimension
        result = contract(tt1, tt2, dims=[1])  # Result has shape (4, 6)
    """
    # Validate inputs
    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}"
        )

    dims_set = set(dims)
    for d in dims:
        if d < 0 or d >= tt1.d:
            raise ValueError(
                f"dim={d} out of range for TT with {tt1.d} dimensions"
            )

    # Special case: contract all dims = dot product
    if dims_set == set(range(tt1.d)):
        matrices = dot(tt1, tt2)
        return matrices[-1].item()

    # Special case: contract no dims = hadamard
    if len(dims) == 0:
        return hadamard(tt1, tt2)

    # General case: hadamard then sum over contracted dims
    # Process dims in reverse order to maintain correct indices
    result = hadamard(tt1, tt2)
    for dim in sorted(dims, reverse=True):
        result = _sum_dim(result, dim)

    return result


def tensordot(
    tt1: TensorTrain,
    tt2: TensorTrain,
    dims: Union[int, tuple],
) -> Union[TensorTrain, float]:
    """
    Compute tensor contraction of two tensor trains (generalized matrix multiply).

    Unlike `contract` which contracts the SAME dimensions and keeps them shared,
    `tensordot` contracts PAIRED dimensions and keeps non-contracted dims from
    BOTH tensors in the result.

    This follows numpy/torch tensordot conventions:
    - dims=n: contract last n dims of tt1 with first n dims of tt2
    - dims=([i,j,...], [k,l,...]): contract tt1's dims i,j,... with tt2's dims k,l,...

    Example for density matrices:
        # psi has shape (n_system, n_env)
        # rho[i,j] = sum_k psi[i,k] * psi[j,k]
        rho = tensordot(psi, psi, dims=([1], [1]))  # shape (n_system, n_system)

    Args:
        tt1: First tensor train
        tt2: Second tensor train
        dims: Contraction specification (int or tuple of lists)

    Returns:
        If all dims contracted: scalar float
        Otherwise: TensorTrain with shape (*kept_dims_tt1, *kept_dims_tt2)

    Note:
        For "last dims" contractions (dims=([d-n,...,d-1], [d-n,...,d-1])) with
        multiple kept dimensions from both tensors, tt2's kept dimensions appear
        in REVERSED order due to TT bond structure constraints. For single kept
        dimensions (common density matrix case), this has no effect.
    """
    # Parse dims argument
    if isinstance(dims, int):
        n = dims
        dims_tt1 = list(range(tt1.d - n, tt1.d))
        dims_tt2 = list(range(n))
    else:
        dims_tt1, dims_tt2 = dims
        dims_tt1 = list(dims_tt1)
        dims_tt2 = list(dims_tt2)
        n = len(dims_tt1)

    # Validate
    if len(dims_tt1) != len(dims_tt2):
        raise ValueError("Must contract same number of dimensions from each tensor")

    for i, j in zip(dims_tt1, dims_tt2):
        if i < 0 or i >= tt1.d:
            raise ValueError(f"tt1 dim {i} out of range")
        if j < 0 or j >= tt2.d:
            raise ValueError(f"tt2 dim {j} out of range")
        if tt1.dims[i] != tt2.dims[j]:
            raise ValueError(
                f"Dimension mismatch: tt1.dims[{i}]={tt1.dims[i]} != tt2.dims[{j}]={tt2.dims[j]}"
            )

    # Special case: no contraction = outer product
    if n == 0:
        return _outer(tt1, tt2)

    # Special case: full contraction = dot product
    if n == tt1.d and n == tt2.d:
        return dot(tt1, tt2)[-1].item()

    # Check if this is a boundary contraction (last of tt1 with first of tt2)
    is_boundary = (
        dims_tt1 == list(range(tt1.d - n, tt1.d)) and
        dims_tt2 == list(range(n))
    )

    if is_boundary:
        return _tensordot_boundary(tt1, tt2, n)
    else:
        # General case: need to permute tensors to make it a boundary contraction
        # For now, support the common case of contracting last dims of both
        # (useful for density matrices)
        is_last_dims = (
            dims_tt1 == list(range(tt1.d - n, tt1.d)) and
            dims_tt2 == list(range(tt2.d - n, tt2.d))
        )
        if is_last_dims:
            return _tensordot_last_dims(tt1, tt2, n)
        else:
            raise NotImplementedError(
                "Only boundary contractions (last of tt1 with first of tt2) "
                "or last-dims contractions are currently supported. "
                f"Got dims_tt1={dims_tt1}, dims_tt2={dims_tt2}"
            )


def _outer(tt1: TensorTrain, tt2: TensorTrain) -> TensorTrain:
    """
    Compute outer product of two tensor trains.

    Result has shape (*tt1.dims, *tt2.dims).
    Since TT boundary ranks are 1, we can simply concatenate the cores.
    """
    result_dims = list(tt1.dims) + list(tt2.dims)
    result_ranks = list(tt1.ranks[1:-1]) + [1] + list(tt2.ranks[1:-1])

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

    for k in range(tt1.d):
        result.cores[k] = tt1.cores[k].clone()
    for k in range(tt2.d):
        result.cores[tt1.d + k] = tt2.cores[k].clone()

    return result


def _tensordot_boundary(tt1: TensorTrain, tt2: TensorTrain, n: int) -> TensorTrain:
    """
    Tensordot where we contract last n dims of tt1 with first n dims of tt2.

    Result has shape (*tt1.dims[:-n], *tt2.dims[n:]).
    """
    # Compute bridge matrix connecting tt1's kept part to tt2's kept part
    # Bridge[i, l] = sum over contracted indices of tt1[..., i, contracted] * tt2[contracted, l, ...]

    # First iteration: A's left index is the external i, B's left index is 1 (boundary)
    A0 = tt1.cores[tt1.d - n]      # (r1_ext, p, r1_next)
    B0 = tt2.cores[0]              # (1, p, r2_next)

    # L[i, j, m] = sum_p A0[i, p, j] * B0[0, p, m]
    L = torch.einsum('ipj,pm->ijm', A0, B0.squeeze(0))

    # Subsequent iterations
    for k in range(1, n):
        A = tt1.cores[tt1.d - n + k]  # (r1_l, p, r1_r)
        B = tt2.cores[k]               # (r2_l, p, r2_r)

        # L[i, j, m] -> L_new[i, j', m']
        # sum_{j, m, p} L[i, j, m] * A[j, p, j'] * B[m, p, m']
        L = torch.einsum('ijm,jpk,mpl->ikl', L, A, B)

    # L has shape (r1_ext, 1, r2_ext) where r1_ext = tt1.ranks[d1-n], r2_ext = tt2.ranks[n]
    Bridge = L.squeeze(1)  # (r1_ext, r2_ext)

    # Build result TT
    kept_dims_tt1 = tt1.dims[:-n] if n < tt1.d else []
    kept_dims_tt2 = tt2.dims[n:] if n < tt2.d else []
    result_dims = list(kept_dims_tt1) + list(kept_dims_tt2)

    if len(kept_dims_tt1) == 0:
        # Only tt2's kept dims
        result_ranks = [tt2.ranks[k] for k in range(n + 1, tt2.d)]
        result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                             tt2.dtype, tt2.device)
        # Absorb Bridge into first kept core of tt2
        first_core = tt2.cores[n]
        result.cores[0] = torch.einsum('ij,jkl->ikl', Bridge, first_core)
        for k in range(n + 1, tt2.d):
            result.cores[k - n] = tt2.cores[k].clone()
        return result

    if len(kept_dims_tt2) == 0:
        # Only tt1's kept dims
        result_ranks = [tt1.ranks[k] for k in range(1, tt1.d - n)]
        result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                             tt1.dtype, tt1.device)
        for k in range(tt1.d - n - 1):
            result.cores[k] = tt1.cores[k].clone()
        # Absorb Bridge into last kept core of tt1
        last_core = tt1.cores[tt1.d - n - 1]
        result.cores[-1] = torch.einsum('ijk,kl->ijl', last_core, Bridge)
        return result

    # Both have kept dims: absorb Bridge by modifying the connection
    # tt1's last kept core connects via Bridge to tt2's first kept core
    result_ranks = (
        [tt1.ranks[k] for k in range(1, tt1.d - n)] +
        [Bridge.shape[1]] +  # connection through Bridge
        [tt2.ranks[k] for k in range(n + 1, tt2.d)]
    )

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

    # Copy tt1's kept cores, absorbing Bridge into the last one
    for k in range(tt1.d - n - 1):
        result.cores[k] = tt1.cores[k].clone()
    last_kept_tt1 = tt1.cores[tt1.d - n - 1]
    result.cores[tt1.d - n - 1] = torch.einsum('ijk,kl->ijl', last_kept_tt1, Bridge)

    # Copy tt2's kept cores
    for k in range(n, tt2.d):
        result.cores[tt1.d - n + k - n] = tt2.cores[k].clone()

    return result


def _tensordot_last_dims(tt1: TensorTrain, tt2: TensorTrain, n: int) -> TensorTrain:
    """
    Tensordot where we contract last n dims of tt1 with last n dims of tt2.

    This is the common case for density matrices:
    rho[i,j] = sum_k psi[i,k] * psi[j,k]

    Result has shape (*tt1.dims[:-n], *reversed(tt2.dims[:-n])).

    Note: Due to TT bond structure constraints, when both tensors have multiple
    kept dimensions, tt2's kept dimensions appear in REVERSED order in the result.
    For a single kept dimension (the common density matrix case), this has no effect.
    For multiple kept dimensions, the user may need to permute the result tensor.
    """
    # Compute bridge connecting tt1's kept part to tt2's kept part
    # Both contracted regions are at the END of their respective TTs

    # For tt1: contracted cores are tt1.d-n to tt1.d-1
    # For tt2: contracted cores are tt2.d-n to tt2.d-1

    # First iteration: both A's and B's left indices are external
    A0 = tt1.cores[tt1.d - n]      # (r1_ext, p, r1_next)
    B0 = tt2.cores[tt2.d - n]      # (r2_ext, p, r2_next)

    # L[i, l, j, m] = sum_p A0[i, p, j] * B0[l, p, m]
    L = torch.einsum('ipj,lpm->iljm', A0, B0)

    # Subsequent iterations
    for k in range(1, n):
        A = tt1.cores[tt1.d - n + k]  # (r1_l, p, r1_r)
        B = tt2.cores[tt2.d - n + k]  # (r2_l, p, r2_r)

        # L[i, l, j, m] -> L_new[i, l, j', m']
        # sum over j, m, p: L[i, l, j, m] * A[j, p, j'] * B[m, p, m']
        L = torch.einsum('iljm,jpk,mpo->ilko', L, A, B)

    # L has shape (r1_ext, r2_ext, 1, 1)
    Bridge = L.squeeze(-1).squeeze(-1)  # (r1_ext, r2_ext)

    # Build result TT with shape (*tt1.dims[:-n], *tt2.dims[:-n] REVERSED)
    # Note: tt2's dims are reversed because we need to flip the TT direction
    # to connect via Bridge. This is a consequence of the TT structure.
    kept_dims_tt1 = tt1.dims[:-n] if n < tt1.d else []
    kept_dims_tt2 = tt2.dims[:-n] if n < tt2.d else []
    # Reverse tt2's dims to match the reversed core order
    result_dims = list(kept_dims_tt1) + list(reversed(kept_dims_tt2))

    d1_kept = len(kept_dims_tt1)
    d2_kept = len(kept_dims_tt2)

    if d1_kept == 0 and d2_kept == 0:
        # Full contraction
        return Bridge.sum().item()

    if d1_kept == 0:
        # Only tt2's kept dims
        result_ranks = [tt2.ranks[k] for k in range(1, tt2.d - n)]
        result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                             tt2.dtype, tt2.device)
        for k in range(d2_kept - 1):
            result.cores[k] = tt2.cores[k].clone()
        # Absorb Bridge into last kept core of tt2
        last_core = tt2.cores[d2_kept - 1]
        result.cores[-1] = torch.einsum('ijk,lk->ijl', last_core,
                                        Bridge.sum(dim=0, keepdim=True))
        return result

    if d2_kept == 0:
        # Only tt1's kept dims
        result_ranks = [tt1.ranks[k] for k in range(1, tt1.d - n)]
        result = TensorTrain(result_dims, result_ranks if result_ranks else None,
                             tt1.dtype, tt1.device)
        for k in range(d1_kept - 1):
            result.cores[k] = tt1.cores[k].clone()
        # Absorb Bridge into last kept core of tt1
        last_core = tt1.cores[d1_kept - 1]
        result.cores[-1] = torch.einsum('ijk,kl->ijl', last_core,
                                        Bridge.sum(dim=1, keepdim=True))
        return result

    # Both have kept dims
    # Result structure: [tt1 kept cores] -- Bridge -- [tt2 kept cores reversed]
    # Bridge connects tt1's last kept core (right bond r1_ext) to
    # tt2's kept region (which ends with right bond r2_ext)

    # For tt2's kept cores, we need to REVERSE their order and TRANSPOSE bond dims
    # because they come from the start of tt2 but need to connect via r2_ext

    # Result ranks: tt1's internal ranks, then Bridge connection, then tt2's reversed ranks
    result_ranks = (
        [tt1.ranks[k] for k in range(1, tt1.d - n)] +
        [Bridge.shape[1]] +
        [tt2.ranks[tt2.d - n - k] for k in range(1, d2_kept)]
    )

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

    # Copy tt1's kept cores, absorb Bridge into last one
    for k in range(d1_kept - 1):
        result.cores[k] = tt1.cores[k].clone()
    last_kept_tt1 = tt1.cores[d1_kept - 1]
    result.cores[d1_kept - 1] = torch.einsum('ijk,kl->ijl', last_kept_tt1, Bridge)

    # Copy tt2's kept cores in REVERSE order with TRANSPOSED bonds
    # This effectively "flips" tt2's TT direction to connect properly
    for k in range(d2_kept):
        orig_k = d2_kept - 1 - k  # reverse order
        core = tt2.cores[orig_k]
        # Transpose bond dimensions: (left, phys, right) -> (right, phys, left)
        result.cores[d1_kept + k] = core.permute(2, 1, 0)

    return result
