
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron GPU Acceleration
#
#  If CuPy is installed and a GPU is available, all matrix math
#  runs on the GPU — much faster for large circuits.
#  If no GPU is found, it automatically falls back to NumPy (CPU).
#
#  You do not need to change your code — just use GPUStatevector
#  or GPUDensityMatrix instead of the normal simulators.
# ─────────────────────────────────────────────────────────────────

def _get_array_module():
    """
    Returns cupy if GPU is available, otherwise numpy.
    This is the core of the fallback system.
    """
    try:
        import cupy as cp
        test = cp.zeros(1)
        test[0] = 1.0  # test item assignment actually works
        return cp, True
    except Exception:
        return np, False


class GPUStatevectorSimulator:
    """
    Statevector simulator with optional GPU acceleration via CuPy.
    Falls back to NumPy automatically if no GPU is available.

    Usage:
        sim = GPUStatevectorSimulator()
        result = sim.run(circuit, shots=1024)
        print(result.get_counts())
        print(f"Used GPU: {sim.using_gpu}")
    """

    def __init__(self):
        self.xp, self.using_gpu = _get_array_module()
        self.name = "gpu_statevector_simulator" if self.using_gpu else "cpu_statevector_simulator"

    def properties(self):
        return {
            "name": self.name,
            "using_gpu": self.using_gpu,
            "backend": "CuPy" if self.using_gpu else "NumPy",
            "simulator": True,
        }

    def run(self, circuit, shots=1024):
        xp = self.xp
        n = circuit.num_qubits

        # Build gate matrices on correct device (GPU or CPU)
        def H():
            return xp.array([[1,1],[1,-1]], dtype=complex) / xp.sqrt(xp.array(2.0))
        def X():
            return xp.array([[0,1],[1,0]], dtype=complex)
        def Y():
            return xp.array([[0,-1j],[1j,0]], dtype=complex)
        def Z():
            return xp.array([[1,0],[0,-1]], dtype=complex)
        def S():
            return xp.array([[1,0],[0,1j]], dtype=complex)
        def T():
            return xp.array([[1,0],[0,xp.exp(1j*xp.pi/4)]], dtype=complex)
        def Rx(t):
            c, s = xp.cos(t/2), xp.sin(t/2)
            return xp.array([[c,-1j*s],[-1j*s,c]], dtype=complex)
        def Ry(t):
            c, s = xp.cos(t/2), xp.sin(t/2)
            return xp.array([[c,-s],[s,c]], dtype=complex)
        def Rz(t):
            return xp.array([[xp.exp(-1j*t/2),0],[0,xp.exp(1j*t/2)]], dtype=complex)

        def full_gate(gate, qubit, num_qubits):
            """Expand single-qubit gate to full n-qubit space."""
            result = xp.array([[1.0+0j]])
            for i in range(num_qubits):
                result = xp.kron(result, gate if i == qubit else xp.eye(2, dtype=complex))
            return result

        # Initialize state |00...0> on GPU or CPU
        state = xp.zeros(2**n, dtype=complex)
        state[0] = 1.0

        # Apply gates
        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]
            gate_mat = None

            try:
                if gname == "H":
                    gate_mat = full_gate(H(), int(parts[-1]), n)
                elif gname == "X":
                    gate_mat = full_gate(X(), int(parts[-1]), n)
                elif gname == "Y":
                    gate_mat = full_gate(Y(), int(parts[-1]), n)
                elif gname == "Z":
                    gate_mat = full_gate(Z(), int(parts[-1]), n)
                elif gname == "S":
                    gate_mat = full_gate(S(), int(parts[-1]), n)
                elif gname == "T":
                    gate_mat = full_gate(T(), int(parts[-1]), n)
                elif gname == "Rx":
                    theta = float(parts[1].strip("()"))
                    gate_mat = full_gate(Rx(theta), int(parts[-1]), n)
                elif gname == "Ry":
                    theta = float(parts[1].strip("()"))
                    gate_mat = full_gate(Ry(theta), int(parts[-1]), n)
                elif gname == "Rz":
                    theta = float(parts[1].strip("()"))
                    gate_mat = full_gate(Rz(theta), int(parts[-1]), n)

                if gate_mat is not None:
                    state = xp.dot(gate_mat, state)
                    norm = xp.sqrt(xp.sum(xp.abs(state)**2))
                    if norm > 0:
                        state /= norm
            except Exception:
                pass

        # Get probabilities — move back to CPU if on GPU
        probs_gpu = xp.abs(state)**2
        probs = np.array(probs_gpu.tolist()) if self.using_gpu else np.array(probs_gpu)
        probs = np.abs(probs) / probs.sum()

        # Sample counts
        counts = {}
        for _ in range(shots):
            outcome = np.random.choice(range(2**n), p=probs)
            key = format(outcome, f"0{n}b")
            counts[key] = counts.get(key, 0) + 1

        return GPUResult(counts, probs, n, shots, self.using_gpu)


class GPUDensityMatrixSimulator:
    """
    Density matrix simulator with optional GPU acceleration via CuPy.
    Falls back to NumPy automatically if no GPU is available.

    Usage:
        sim = GPUDensityMatrixSimulator()
        result = sim.run(circuit, shots=1024)
        print(result.get_density_matrix())
        print(f"Used GPU: {sim.using_gpu}")
    """

    def __init__(self):
        self.xp, self.using_gpu = _get_array_module()
        self.name = "gpu_density_matrix_simulator" if self.using_gpu else "cpu_density_matrix_simulator"

    def properties(self):
        return {
            "name": self.name,
            "using_gpu": self.using_gpu,
            "backend": "CuPy" if self.using_gpu else "NumPy",
            "simulator": True,
        }

    def run(self, circuit, shots=1024):
        xp = self.xp
        n = circuit.num_qubits
        dim = 2**n

        def full_gate(gate, qubit):
            result = xp.array([[1.0+0j]])
            for i in range(n):
                result = xp.kron(result, gate if i == qubit else xp.eye(2, dtype=complex))
            return result

        # Start in |0><0|
        rho = xp.zeros((dim, dim), dtype=complex)
        rho[0, 0] = 1.0

        def H():
            return xp.array([[1,1],[1,-1]], dtype=complex) / xp.sqrt(xp.array(2.0))
        def X():
            return xp.array([[0,1],[1,0]], dtype=complex)
        def Y():
            return xp.array([[0,-1j],[1j,0]], dtype=complex)
        def Z():
            return xp.array([[1,0],[0,-1]], dtype=complex)
        def Rx(t):
            c, s = xp.cos(t/2), xp.sin(t/2)
            return xp.array([[c,-1j*s],[-1j*s,c]], dtype=complex)
        def Ry(t):
            c, s = xp.cos(t/2), xp.sin(t/2)
            return xp.array([[c,-s],[s,c]], dtype=complex)
        def Rz(t):
            return xp.array([[xp.exp(-1j*t/2),0],[0,xp.exp(1j*t/2)]], dtype=complex)

        # Apply gates as rho -> U rho U†
        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]
            U = None
            try:
                if gname == "H":
                    U = full_gate(H(), int(parts[-1]))
                elif gname == "X":
                    U = full_gate(X(), int(parts[-1]))
                elif gname == "Y":
                    U = full_gate(Y(), int(parts[-1]))
                elif gname == "Z":
                    U = full_gate(Z(), int(parts[-1]))
                elif gname == "Rx":
                    theta = float(parts[1].strip("()"))
                    U = full_gate(Rx(theta), int(parts[-1]))
                elif gname == "Ry":
                    theta = float(parts[1].strip("()"))
                    U = full_gate(Ry(theta), int(parts[-1]))
                elif gname == "Rz":
                    theta = float(parts[1].strip("()"))
                    U = full_gate(Rz(theta), int(parts[-1]))
                if U is not None:
                    rho = U @ rho @ U.conj().T
            except Exception:
                pass

        # Move density matrix back to CPU
        rho_cpu = np.array(rho.tolist()) if self.using_gpu else np.array(rho)

        # Sample counts
        probs = np.abs(np.real(np.diag(rho_cpu)))
        probs /= probs.sum()
        counts = {}
        for _ in range(shots):
            outcome = np.random.choice(range(dim), p=probs)
            key = format(outcome, f"0{n}b")
            counts[key] = counts.get(key, 0) + 1

        return GPUResult(counts, probs, n, shots, self.using_gpu, rho_cpu)


class GPUResult:
    """Holds results from GPU or CPU simulation."""
    def __init__(self, counts, probabilities, num_qubits, shots, used_gpu, density_matrix=None):
        self.counts = counts
        self.probabilities = probabilities
        self.num_qubits = num_qubits
        self.shots = shots
        self.used_gpu = used_gpu
        self._density_matrix = density_matrix

    def get_counts(self):
        return self.counts

    def get_probabilities(self):
        return self.probabilities

    def get_density_matrix(self):
        if self._density_matrix is None:
            raise ValueError("No density matrix. Use GPUDensityMatrixSimulator.")
        return self._density_matrix

    def __repr__(self):
        backend = "GPU" if self.used_gpu else "CPU"
        return (f"GPUResult(backend={backend}, shots={self.shots}, "
                f"counts={self.counts})")
