
import numpy as np
import re

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

class HPDensityMatrixResult:
    def __init__(self, counts, rho, n):
        self._counts = counts
        self._rho = rho
        self.num_qubits = n
    def get_counts(self): return self._counts
    def get_density_matrix(self): return self._rho
    def purity(self): return float(np.real(np.trace(self._rho @ self._rho)))
    def entropy(self):
        vals = np.linalg.eigvalsh(self._rho)
        vals = vals[vals > 1e-12]
        return float(-np.sum(vals * np.log2(vals)))
    def trace(self): return float(np.real(np.trace(self._rho)))
    def __repr__(self):
        return (f"HPDensityMatrixResult(qubits={self.num_qubits}, "
                f"purity={self.purity():.4f}, trace={self.trace():.4f})")


class HPDensityMatrixSimulator:
    _I    = np.eye(2, dtype=complex)
    _H    = np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)
    _X    = np.array([[0,1],[1,0]], dtype=complex)
    _Y    = np.array([[0,-1j],[1j,0]], dtype=complex)
    _Z    = np.array([[1,0],[0,-1]], dtype=complex)
    _S    = np.array([[1,0],[0,1j]], dtype=complex)
    _T    = np.array([[1,0],[0,np.exp(1j*np.pi/4)]], dtype=complex)
    _CNOT = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]], dtype=complex)
    _CZ   = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,-1]], dtype=complex)
    _SWAP = np.array([[1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]], dtype=complex)

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

    def _to_xp(self, arr):
        return self.xp.array(arr) if self.use_gpu else np.array(arr)

    def _to_np(self, arr):
        return cp.asnumpy(arr) if (self.use_gpu and HAS_GPU) else np.array(arr)

    def _zero_rho(self, n):
        xp = self.xp
        rho = xp.zeros((2**n, 2**n), dtype=complex)
        rho[0, 0] = 1.0 + 0j
        return rho

    def _expand_1q(self, U, qubit, n):
        """Expand a 2x2 gate to the full 2^n x 2^n unitary."""
        full = None
        for i in range(n):
            m = U if i == qubit else self._I
            full = m if full is None else np.kron(full, m)
        return self._to_xp(full)

    def _expand_2q(self, U4, q0, q1, n):
        """Build full 2^n x 2^n unitary for a 2-qubit gate on q0, q1."""
        dim = 2**n
        result = np.zeros((dim, dim), dtype=complex)
        for col in range(dim):
            bits = [(col >> (n-1-i)) & 1 for i in range(n)]
            b0, b1 = bits[q0], bits[q1]
            amp_in = np.zeros(4, dtype=complex)
            amp_in[b0*2+b1] = 1.0
            amp_out = U4 @ amp_in
            for out_idx in range(4):
                ob0, ob1 = out_idx >> 1, out_idx & 1
                new_bits = bits.copy()
                new_bits[q0] = ob0; new_bits[q1] = ob1
                row = sum(new_bits[i] << (n-1-i) for i in range(n))
                result[row, col] += amp_out[out_idx]
        return self._to_xp(result)

    def _apply_unitary(self, rho, U):
        return U @ rho @ self.xp.conj(U).T

    def _apply_kraus(self, rho, kraus_list, qubit, n):
        """Apply single-qubit Kraus ops expanded to full system size."""
        xp = self.xp
        result = xp.zeros_like(rho)
        for K in kraus_list:
            # Expand 2x2 K to full 2^n space
            K_full = None
            for i in range(n):
                m = K if i == qubit else self._I
                K_full = m if K_full is None else np.kron(K_full, m)
            K_xp = self._to_xp(K_full)
            result = result + K_xp @ rho @ xp.conj(K_xp).T
        return result

    def _get_gate(self, name, params):
        if name == "H":  return self._H
        if name == "X":  return self._X
        if name == "Y":  return self._Y
        if name == "Z":  return self._Z
        if name == "S":  return self._S
        if name == "T":  return self._T
        if name == "RX":
            th = float(params[0]) if params else 0.0
            c,s = np.cos(th/2), np.sin(th/2)
            return np.array([[c,-1j*s],[-1j*s,c]], dtype=complex)
        if name == "RY":
            th = float(params[0]) if params else 0.0
            c,s = np.cos(th/2), np.sin(th/2)
            return np.array([[c,-s],[s,c]], dtype=complex)
        if name == "RZ":
            th = float(params[0]) if params else 0.0
            return np.array([[np.exp(-1j*th/2),0],[0,np.exp(1j*th/2)]], dtype=complex)
        return None

    def run(self, circuit, shots=1024, noise_model=None):
        n = circuit.num_qubits
        rho = self._zero_rho(n)

        for entry in circuit.gate_history:
            e = entry.lower()
            gate_name = None
            for g in ["rx","ry","rz","cnot","cz","swap","h","x","y","z","s","t"]:
                if e.startswith(g):
                    gate_name = g.upper(); break
            if gate_name is None: continue

            nums = list(map(int, re.findall(r"\d+", entry)))
            if not nums: continue
            floats = list(map(float, re.findall(r"[-+]?\d*\.\d+", entry)))

            if gate_name in ("CNOT","CZ","SWAP") and len(nums) >= 2:
                U4 = {"CNOT":self._CNOT,"CZ":self._CZ,"SWAP":self._SWAP}[gate_name]
                U_full = self._expand_2q(U4, nums[0], nums[1], n)
                rho = self._apply_unitary(rho, U_full)
            else:
                U1 = self._get_gate(gate_name, floats)
                if U1 is not None:
                    U_full = self._expand_1q(U1, nums[0], n)
                    rho = self._apply_unitary(rho, U_full)
                    if noise_model is not None:
                        err = noise_model.get_error(gate_name, nums[0])
                        if err:
                            # Pass qubit index and n so Kraus ops are expanded correctly
                            rho = self._apply_kraus(rho, err.kraus(), nums[0], n)

        rho_np = self._to_np(rho)
        probs = np.real(np.diag(rho_np))
        probs = np.clip(probs, 0, None)
        probs /= probs.sum()
        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))}
        return HPDensityMatrixResult(counts, rho_np, n)

    def partial_trace(self, rho_np, n, keep_qubits):
        trace_q = [q for q in range(n) if q not in keep_qubits]
        dm = rho_np.reshape([2]*(2*n))
        for q in sorted(trace_q, reverse=True):
            dm = np.trace(dm, axis1=q, axis2=q+len(dm.shape)//2)
        k = len(keep_qubits)
        return dm.reshape(2**k, 2**k)

    def properties(self):
        return {"name": "HPDensityMatrixSimulator",
                "backend": "CuPy-GPU" if self.use_gpu else "NumPy-CPU",
                "supports_noise": True, "supports_2qubit_gates": True,
                "max_qubits_recommended": 12}
