# Ensure typing is initialized early to avoid numpy/typing circular import issues
try:
    import typing as _typing
except Exception:
    _typing = None

# Try importing numpy safely; perform a couple of retries to avoid partial initialization races
_NUMPY_AVAILABLE = False
np = None
for attempt in range(2):
    try:
        import numpy as np
        import math
        _NUMPY_AVAILABLE = True
        break
    except Exception as e:
        # small fallback: collect garbage and retry once
        import gc
        gc.collect()
        np = None
        _NUMPY_AVAILABLE = False
        last_import_error = e
        if attempt == 1:
            # on final failure, log and proceed with fallback mode
            print(f"WARNING: numpy import failed in mps_utils after retries: {e}")

# Minimal, lossless Tensor-Train (TT) / Matrix Product State (MPS) decomposition
# If numpy is available, we use TT-SVD (lossless if no truncation). Otherwise we provide a simple single-core fallback.

def bytes_to_vector(data: bytes):
    if not data:
        return []
    if _NUMPY_AVAILABLE:
        return np.frombuffer(data, dtype=np.uint8).astype(np.float64)
    else:
        return [b for b in data]


def _compute_dims(n, mode_size=256):
    # Split length into dims where most dims = mode_size and last may be smaller
    if n == 0:
        return []
    L = math.ceil(n / mode_size)
    dims = [mode_size] * (L - 1)
    last = n - mode_size * (L - 1)
    dims.append(last)
    return dims


def tt_decompose(data: bytes, mode_size: int = 256):
    """Decompose byte sequence into a list of TT cores.

    If numpy is not available, return a single-core representation storing raw bytes.
    """
    if not data:
        return []
    if not _NUMPY_AVAILABLE:
        # single-core fallback: store as raw bytes in a single element list
        return [data]

    vec = bytes_to_vector(data)
    n = vec.size
    if n == 0:
        return []
    dims = _compute_dims(n, mode_size)
    tensor = vec.reshape(dims)

    cores = []
    cur = tensor
    r = 1
    for k in range(len(dims) - 1):
        dk = dims[k]
        cur_shape = (r * dk, -1)
        mat = cur.reshape(cur_shape)
        # full SVD
        U, S, Vt = np.linalg.svd(mat, full_matrices=False)
        r_new = U.shape[1]
        core = U.reshape(r, dk, r_new)
        cores.append(core)
        # build next cur = diag(S) @ Vt
        cur = (np.diag(S) @ Vt).reshape(r_new, *dims[k + 1:])
        r = r_new
    # last core
    last_core = cur.reshape(r, dims[-1], 1)
    cores.append(last_core)
    return cores


def tt_reconstruct(cores):
    """Reconstruct original byte sequence from TT cores."""
    if not cores:
        return b""
    if not _NUMPY_AVAILABLE:
        # single-core fallback
        return cores[0]
    cur = cores[0]
    for core in cores[1:]:
        # contract the right index of cur with left index of core
        cur = np.tensordot(cur, core, axes=([2], [0]))
    vec = cur.reshape(-1)
    # Round and clip to uint8 to recover original bytes (should be exact if no truncation)
    vec = np.rint(vec).astype(np.int64)
    vec = np.clip(vec, 0, 255).astype(np.uint8)
    return vec.tobytes()


# Small sanity check function
def roundtrip_bytes(data: bytes, mode_size: int = 256):
    cores = tt_decompose(data, mode_size=mode_size)
    return tt_reconstruct(cores)
