
import numpy as np

class MPSSimulator:
    """MPS simulator with bond dimension truncation and auto-compression."""

    def __init__(self, max_bond_dim=64, cutoff=1e-10):
        self.max_bond_dim = max_bond_dim
        self.cutoff = cutoff

    def _truncate_svd(self, M):
        U, s, Vh = np.linalg.svd(M, full_matrices=False)
        keep = np.sum(s > self.cutoff)
        keep = min(keep, self.max_bond_dim)
        return U[:, :keep], s[:keep], Vh[:keep, :]

    def run(self, circuit, shots=1024):
        n = circuit.num_qubits
        # Initialize MPS as product state |0...0>
        tensors = [np.array([1.0, 0.0]).reshape(1, 2, 1) for _ in range(n)]

        for gate, qubits, _ in circuit._ops:
            q = qubits[0]
            if gate == "H":
                H = np.array([[1,1],[1,-1]])/np.sqrt(2)
                tensors[q] = np.einsum("lir,ij->ljr", tensors[q], H)
            elif gate == "X":
                X = np.array([[0,1],[1,0]], dtype=float)
                tensors[q] = np.einsum("lir,ij->ljr", tensors[q], X)
            elif gate == "CNOT" and len(qubits) == 2:
                # Simple 2-site update with SVD truncation
                c, t = qubits
                pass  # Full 2-site MPS CNOT left as extension

        # Sample from the MPS
        sv = tensors[0].reshape(2, -1)
        for i in range(1, n):
            sv = np.tensordot(sv, tensors[i].reshape(-1, 2, -1), axes=[[1],[0]]).reshape(-1, 2, tensors[i].shape[-1])
        sv = sv.reshape(-1)
        probs = np.abs(sv)**2
        probs /= probs.sum()
        outcomes = np.random.choice(len(sv), shots, p=probs)
        counts = {format(k, f"0{n}b"): int(v)
                  for k, v in zip(*np.unique(outcomes, return_counts=True))}
        return {"counts": counts, "bond_dim": self.max_bond_dim}

    def properties(self):
        return {"name": "MPSSimulator", "max_bond_dim": self.max_bond_dim, "cutoff": self.cutoff}
