
import numpy as np
import re

try:
    import cupy as cp
    _t = cp.zeros(1)
    _t[0] = 1.0
    HAS_GPU = True
except Exception:
    cp = None
    HAS_GPU = False

def _xp():
    if not HAS_GPU:
        return np
    try:
        test = cp.zeros(1)
        test[0] = 1.0
        return cp
    except Exception:
        return np

def _to_np(arr):
    if HAS_GPU and cp is not None and isinstance(arr, cp.ndarray):
        return cp.asnumpy(arr)
    return np.array(arr)

# ── Gate matrices ────────────────────────────────────────────────
_SQ2 = np.sqrt(2.0)
GATE_LIB_NP = {
    "H":  np.array([[1,1],[1,-1]], dtype=complex) / _SQ2,
    "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),
}
TWO_QUBIT_NP = {
    "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),
    "CY":   np.array([[1,0,0,0],[0,1,0,0],[0,0,0,-1j],[0,0,1j,0]], dtype=complex),
    "SWAP": np.array([[1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]], dtype=complex),
}

# ── Parse gate_history into structured ops ────────────────────────
# Exact formats from gates.py:
#   "H on qubit 0"
#   "Rx(1.57) on qubit 0"
#   "CNOT control 0 target 1"
#   "CZ control 0 target 1"
#   "CY control 0 target 1"
#   "SWAP qubits 0 and 1"
#   "Toffoli control1 0 control2 1 target 2"

def _parse_history(circuit):
    ops = []
    for entry in circuit.gate_history:
        e = entry.strip()
        el = e.lower()

        # Two-qubit: CNOT / CZ / CY
        m = re.match(r"(cnot|cz|cy)\s+control\s+(\d+)\s+target\s+(\d+)", el)
        if m:
            ops.append((m.group(1).upper(), [int(m.group(2)), int(m.group(3))], []))
            continue

        # SWAP
        m = re.match(r"swap\s+qubits\s+(\d+)\s+and\s+(\d+)", el)
        if m:
            ops.append(("SWAP", [int(m.group(1)), int(m.group(2))], []))
            continue

        # Toffoli
        m = re.match(r"toffoli\s+control1\s+(\d+)\s+control2\s+(\d+)\s+target\s+(\d+)", el)
        if m:
            ops.append(("TOFFOLI", [int(m.group(1)), int(m.group(2)), int(m.group(3))], []))
            continue

        # Parameterized: Rx/Ry/Rz
        m = re.match(r"r([xyz])\(([-\d.]+)\)\s+on\s+qubit\s+(\d+)", el)
        if m:
            ops.append(("R"+m.group(1).upper(), [int(m.group(3))], [float(m.group(2))]))
            continue

        # Single-qubit no param: "H on qubit 0"
        m = re.match(r"([a-z†]+)\s+on\s+qubit\s+(\d+)", el)
        if m:
            gname = m.group(1).upper().replace("†","DG")
            ops.append((gname, [int(m.group(2))], []))
            continue

    return ops

# ── Apply single-qubit gate to statevector ────────────────────────
def _apply_1q(xp, sv, G, qubit, n):
    G_xp = xp.array(G, dtype=complex)
    sv = sv.reshape([2]*n)
    # contract: new[...P...] = G[P,p] * sv[...p...]
    sv = xp.tensordot(G_xp, sv, axes=[[1],[qubit]])
    # tensordot puts the new axis at position 0; move it back to qubit
    order = list(range(1, qubit+1)) + [0] + list(range(qubit+1, n))
    sv = xp.transpose(sv, order)
    return sv.reshape(2**n)

# ── Apply two-qubit gate to statevector ───────────────────────────
def _apply_2q(xp, sv, G4, q0, q1, n):
    G_xp = xp.array(G4.reshape(2,2,2,2), dtype=complex)
    sv = sv.reshape([2]*n)
    # contract over axes q0,q1 simultaneously
    sv = xp.tensordot(G_xp, sv, axes=[[2,3],[q0,q1]])
    # result shape: (2,2, ...n-2 axes...)
    # axes 0,1 correspond to output q0,q1; rest are in original order
    # build permutation to put them back
    remaining = [i for i in range(n) if i not in (q0,q1)]
    # current order: [out_q0=0, out_q1=1, remaining...]
    # target order: insert 0->q0, 1->q1 into remaining positions
    full = list(range(2, n))   # remaining axes (0-indexed after the 2 output axes)
    perm = [None] * n
    perm[q0] = 0
    perm[q1] = 1
    ri = 2
    for pos in remaining:
        perm[pos] = ri; ri += 1
    sv = xp.transpose(sv, perm)
    return sv.reshape(2**n)

# ── Build rotation gates ──────────────────────────────────────────
def _rot_gate(axis, theta):
    c, s = np.cos(theta/2), np.sin(theta/2)
    if axis == "X": return np.array([[c,-1j*s],[-1j*s,c]], dtype=complex)
    if axis == "Y": return np.array([[c,-s],[s,c]], dtype=complex)
    if axis == "Z": return np.array([[np.exp(-1j*theta/2),0],[0,np.exp(1j*theta/2)]], dtype=complex)

# ── Run one circuit, return probability vector (on device) ────────
def _run_sv_probs(xp, circuit):
    n = circuit.num_qubits
    sv = xp.zeros(2**n, dtype=complex)
    sv[0] = 1.0 + 0j

    ops = _parse_history(circuit)
    for gate, qubits, params in ops:
        if gate in GATE_LIB_NP:
            sv = _apply_1q(xp, sv, GATE_LIB_NP[gate], qubits[0], n)
        elif gate in ("RX","RY","RZ"):
            th = params[0] if params else 0.0
            G = _rot_gate(gate[1], th)
            sv = _apply_1q(xp, sv, G, qubits[0], n)
        elif gate in TWO_QUBIT_NP and len(qubits) >= 2:
            sv = _apply_2q(xp, sv, TWO_QUBIT_NP[gate], qubits[0], qubits[1], n)
        # Toffoli and unknown gates: skip (rare in basic circuits)

    probs = xp.abs(sv)**2
    probs = probs / probs.sum()
    return probs

# ── GPUBatchedSampler ─────────────────────────────────────────────
class GPUBatchedSampler:
    """
    Run many circuits in a batch — all on GPU (or NumPy fallback).

    Usage:
        sampler = GPUBatchedSampler(shots=1024)
        result  = sampler.run(qc)
        results = sampler.run_batch([qc1, qc2, qc3])
    """
    def __init__(self, shots=1024):
        self.shots = shots
        self.xp = _xp()
        self.using_gpu = HAS_GPU

    def run(self, circuit, shots=None):
        shots = shots or self.shots
        xp = self.xp
        n = circuit.num_qubits
        probs_dev = _run_sv_probs(xp, circuit)
        probs_np = _to_np(probs_dev)
        probs_np = np.clip(probs_np, 0, None)
        probs_np /= probs_np.sum()
        outcomes = np.random.choice(2**n, shots, p=probs_np)
        counts = {format(k, f"0{n}b"): int(v)
                  for k, v in zip(*np.unique(outcomes, return_counts=True))}
        return {"counts": counts, "using_gpu": self.using_gpu}

    def run_batch(self, circuits, shots=None):
        return [self.run(qc, shots) for qc in circuits]

    def properties(self):
        return {"name": "GPUBatchedSampler",
                "backend": "CuPy-GPU" if self.using_gpu else "NumPy-CPU",
                "shots": self.shots}


# ── GPUBatchedEstimator ───────────────────────────────────────────
_PAULIS = {
    "I": np.eye(2,dtype=complex),
    "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),
}

class GPUBatchedEstimator:
    """
    Compute expectation values of Pauli strings on GPU.

    Usage:
        est = GPUBatchedEstimator()
        r   = est.run(qc, "ZZ")
        rs  = est.run_batch([(qc1,"ZZ"),(qc2,"XI")])
    """
    def __init__(self):
        self.xp = _xp()
        self.using_gpu = HAS_GPU

    def _build_obs(self, pauli_str, n):
        mat = np.array([[1.0+0j]])
        ps = pauli_str.upper().ljust(n,"I")[:n]
        for c in ps:
            mat = np.kron(mat, _PAULIS.get(c, _PAULIS["I"]))
        return self.xp.array(mat, dtype=complex)

    def run(self, circuit, pauli_str):
        xp = self.xp
        n = circuit.num_qubits
        sv = xp.zeros(2**n, dtype=complex); sv[0] = 1.0+0j
        ops = _parse_history(circuit)
        for gate, qubits, params in ops:
            if gate in GATE_LIB_NP:
                sv = _apply_1q(xp, sv, GATE_LIB_NP[gate], qubits[0], n)
            elif gate in ("RX","RY","RZ"):
                sv = _apply_1q(xp, sv, _rot_gate(gate[1], params[0] if params else 0.0), qubits[0], n)
            elif gate in TWO_QUBIT_NP and len(qubits) >= 2:
                sv = _apply_2q(xp, sv, TWO_QUBIT_NP[gate], qubits[0], qubits[1], n)
        O = self._build_obs(pauli_str, n)
        exp_val = float(xp.real(xp.conj(sv) @ O @ sv))
        return {"expectation_value": exp_val, "pauli": pauli_str,
                "using_gpu": self.using_gpu}

    def run_batch(self, jobs):
        return [self.run(qc, obs) for qc, obs in jobs]


# ── Convenience one-liner ─────────────────────────────────────────
def gpu_execute(circuit, shots=1024, observable=None):
    """
    One-line GPU execution.
        counts = gpu_execute(qc, shots=1024)
        exp    = gpu_execute(qc, observable="ZZ")
    """
    if observable:
        return GPUBatchedEstimator().run(circuit, observable)
    return GPUBatchedSampler(shots=shots).run(circuit)
