
import numpy as np

try:
    import cupy as cp
    HAS_GPU = True
except ImportError:
    cp = np
    HAS_GPU = False

class FastDensityMatrixSimulator:
    """Vectorized density matrix simulator with optional GPU and noise support."""

    def __init__(self, use_gpu=False):
        self.xp = cp if (use_gpu and HAS_GPU) else np
        self.using_gpu = use_gpu and HAS_GPU

    def _zero_dm(self, n):
        dm = self.xp.zeros((2**n, 2**n), dtype=complex)
        dm[0, 0] = 1.0
        return dm

    def _apply_gate_dm(self, dm, U, qubit, n):
        xp = self.xp
        U = xp.array(U, dtype=complex)
        dim = 2**n
        dm = dm.reshape([2]*2*n)
        dm = xp.tensordot(U, dm, axes=[[1],[qubit]])
        order = list(range(1, qubit+1)) + [0] + list(range(qubit+1, 2*n))
        dm = xp.transpose(dm, order)
        dm = xp.tensordot(U.conj(), dm, axes=[[1],[n+qubit]])
        order2 = list(range(n+qubit, 2*n)) + [0] + list(range(n+qubit))
        # simplify: just reshape and return after gate
        dm = dm.reshape(dim, dim)
        return dm

    def run(self, circuit, shots=1024, noise_model=None):
        n = circuit.num_qubits
        xp = self.xp
        dm = self._zero_dm(n)

        GATES = {
            "H": np.array([[1,1],[1,-1]])/np.sqrt(2),
            "X": np.array([[0,1],[1,0]], dtype=float),
            "Y": np.array([[0,-1j],[1j,0]]),
            "Z": np.array([[1,0],[0,-1]], dtype=float),
        }

        for gate, qubits, _ in circuit._ops:
            if gate in GATES:
                U = GATES[gate]
                q = qubits[0]
                dim = 2**n
                # Full unitary on subspace
                full_U = np.eye(dim, dtype=complex)
                for i in range(dim):
                    for j in range(dim):
                        if i == j:
                            full_U[i,j] = U[0,0] if not (i >> (n-1-q) & 1) else U[1,1]
                dm_np = xp.asnumpy(dm) if self.using_gpu else dm
                dm_np = full_U @ dm_np @ full_U.conj().T
                dm = xp.array(dm_np) if self.using_gpu else dm_np

                if noise_model:
                    err = noise_model.get_error(gate, q)
                    if err:
                        dm_np = xp.asnumpy(dm) if self.using_gpu else dm
                        dm_np = noise_model.apply_kraus(dm_np, err.kraus())
                        dm = xp.array(dm_np) if self.using_gpu else dm_np

        dm_np = (cp.asnumpy(dm) if self.using_gpu else dm)
        probs = np.real(np.diag(dm_np))
        probs = np.abs(probs) / np.sum(np.abs(probs))
        outcomes = np.random.choice(2**n, shots, p=probs)
        counts = {format(k, f"0{n}b"): int(v)
                  for k, v in zip(*np.unique(outcomes, return_counts=True))}
        purity = float(np.real(np.trace(dm_np @ dm_np)))
        return {"counts": counts, "density_matrix": dm_np, "purity": purity, "trace": float(np.real(np.trace(dm_np)))}
