
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Quantum Debugger
#
#  Works exactly like a code debugger (like in VS Code or PyCharm)
#  but for quantum circuits.
#
#  You set "breakpoints" at specific gate steps.
#  When you run the circuit, it PAUSES at each breakpoint and
#  shows you the full quantum state at that exact moment:
#    - Which basis states are active
#    - Their probabilities
#    - Their complex amplitudes
#    - Their phase angles
#
#  Usage:
#      dbg = QuantumDebugger()
#      dbg.set_breakpoint(1)   # pause after gate 1
#      dbg.set_breakpoint(3)   # pause after gate 3
#      snapshots = dbg.run_debug(my_circuit)
# ─────────────────────────────────────────────────────────────────

# ── Gate matrices (same as step_simulator) ────────────────────────
_SQ2 = np.sqrt(2.0)

def _H():   return np.array([[1,1],[1,-1]], dtype=complex) / _SQ2
def _X():   return np.array([[0,1],[1,0]], dtype=complex)
def _Y():   return np.array([[0,-1j],[1j,0]], dtype=complex)
def _Z():   return np.array([[1,0],[0,-1]], dtype=complex)
def _S():   return np.array([[1,0],[0,1j]], dtype=complex)
def _Sdg(): return np.array([[1,0],[0,-1j]], dtype=complex)
def _T():   return np.array([[1,0],[0,np.exp(1j*np.pi/4)]], dtype=complex)
def _Tdg(): return np.array([[1,0],[0,np.exp(-1j*np.pi/4)]], dtype=complex)
def _Rx(t):
    c,s = np.cos(t/2), np.sin(t/2)
    return np.array([[c,-1j*s],[-1j*s,c]], dtype=complex)
def _Ry(t):
    c,s = np.cos(t/2), np.sin(t/2)
    return np.array([[c,-s],[s,c]], dtype=complex)
def _Rz(t):
    return np.array([[np.exp(-1j*t/2),0],[0,np.exp(1j*t/2)]], dtype=complex)
def _CNOT(): return np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]], dtype=complex)
def _CZ():   return np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,-1]], dtype=complex)
def _CY():   return np.array([[1,0,0,0],[0,1,0,0],[0,0,0,-1j],[0,0,1j,0]], dtype=complex)
def _SWAP(): return np.array([[1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]], dtype=complex)
def _TOFF():
    g = np.eye(8, dtype=complex)
    g[6,6]=g[7,7]=0; g[6,7]=g[7,6]=1
    return g

# ── Gate application helpers ──────────────────────────────────────

def _apply_1q(sv, mat, q, n):
    psi = sv.reshape([2]*n)
    psi = np.tensordot(mat, psi, axes=[[1],[q]])
    order = list(range(1,q+1)) + [0] + list(range(q+1,n))
    return np.transpose(psi, order).reshape(-1)

def _apply_2q(sv, mat, q0, q1, n):
    psi = sv.reshape([2]*n)
    G   = mat.reshape(2,2,2,2)
    ii  = list(range(n)); oi = list(range(n))
    ni,nj = n, n+1
    oi[q0]=ni; oi[q1]=nj
    return np.einsum(G,[ni,nj,q0,q1], psi,ii, oi).reshape(-1)

def _apply_3q(sv, mat, q0, q1, q2, n):
    psi = sv.reshape([2]*n)
    G   = mat.reshape(2,2,2,2,2,2)
    ii  = list(range(n)); oi = list(range(n))
    ni,nj,nk = n, n+1, n+2
    oi[q0]=ni; oi[q1]=nj; oi[q2]=nk
    return np.einsum(G,[ni,nj,nk,q0,q1,q2], psi,ii, oi).reshape(-1)

# ── Parser (same format as gate_history strings) ──────────────────

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

        m = re.match(r'r([xyz])\(([^)]+)\) on qubit (\d+)', el)
        if m:
            ops.append(('r'+m.group(1), int(m.group(3)), float(m.group(2)))); continue

        m = re.match(r'([a-z†]+) on qubit (\d+)', el)
        if m:
            ops.append((m.group(1).replace('†','dg'), int(m.group(2)))); continue

        m = re.match(r'cnot control (\d+) target (\d+)', el)
        if m:
            ops.append(('cnot', int(m.group(1)), int(m.group(2)))); continue

        m = re.match(r'c([zy]) control (\d+) target (\d+)', el)
        if m:
            ops.append(('c'+m.group(1), int(m.group(2)), int(m.group(3)))); continue

        m = re.match(r'swap qubits (\d+) and (\d+)', el)
        if m:
            ops.append(('swap', int(m.group(1)), int(m.group(2)))); continue

        m = re.match(r'toffoli control1 (\d+) control2 (\d+) target (\d+)', el)
        if m:
            ops.append(('toffoli', int(m.group(1)), int(m.group(2)), int(m.group(3)))); continue

        if 's†' in el or 'sdg' in el:
            m = re.search(r'qubit (\d+)', el)
            if m: ops.append(('sdg', int(m.group(1)))); continue
        if 't†' in el or 'tdg' in el:
            m = re.search(r'qubit (\d+)', el)
            if m: ops.append(('tdg', int(m.group(1)))); continue
    return ops


# ── Main Debugger Class ───────────────────────────────────────────

class QuantumDebugger:
    """
    Quantum circuit debugger with breakpoints.

    Step numbers are 1-based:
      Step 1 = after the 1st gate
      Step 2 = after the 2nd gate
      ...and so on

    Example:
        dbg = QuantumDebugger()
        dbg.set_breakpoint(1)
        dbg.set_breakpoint(2)
        snapshots = dbg.run_debug(my_circuit)

        # Get a specific snapshot later
        snap = dbg.get_snapshot(1)
    """

    def __init__(self):
        self._breakpoints = set()
        self._snapshots   = {}     # step → snapshot dict

    # ── Breakpoint control ────────────────────────────────────────

    def set_breakpoint(self, step_index: int):
        """
        Add a breakpoint at the given step (1-based).
        The debugger will pause and print the state after that gate.
        """
        self._breakpoints.add(step_index)
        print(f"[Debugger] ✅ Breakpoint set at step {step_index}")

    def clear_breakpoint(self, step_index: int):
        """Remove a single breakpoint."""
        self._breakpoints.discard(step_index)
        print(f"[Debugger] Breakpoint at step {step_index} removed")

    def clear_all_breakpoints(self):
        """Remove every breakpoint."""
        self._breakpoints.clear()
        print("[Debugger] All breakpoints cleared")

    def list_breakpoints(self):
        """Show all active breakpoints."""
        if self._breakpoints:
            print(f"[Debugger] Active breakpoints: {sorted(self._breakpoints)}")
        else:
            print("[Debugger] No breakpoints set")

    # ── Run ───────────────────────────────────────────────────────

    def run_debug(self, circuit) -> dict:
        """
        Run the circuit and pause at every breakpoint.

        Prints the quantum state at each breakpoint automatically.

        Returns:
            dict of { step_number -> snapshot_dict }
            Each snapshot has: step, gate_label, statevector,
                               probs, phases, labels
        """
        n   = circuit.num_qubits
        raw = getattr(circuit, 'gate_history', [])
        ops = _parse_history(raw)

        labels = [f'|{format(i, f"0{n}b")}>' for i in range(2**n)]
        sv     = np.zeros(2**n, dtype=complex)
        sv[0]  = 1.0

        self._snapshots = {}

        _1Q = {'h':_H(),'x':_X(),'y':_Y(),'z':_Z(),
               's':_S(),'sdg':_Sdg(),'t':_T(),'tdg':_Tdg()}

        print("=" * 60)
        print("  PKTron Quantum Debugger — Running")
        print(f"  {len(ops)} gate(s) | "
              f"{len(self._breakpoints)} breakpoint(s): "
              f"{sorted(self._breakpoints)}")
        print("=" * 60)

        for i, op in enumerate(ops):
            step = i + 1          # 1-based
            gate = op[0].lower()

            # Apply gate
            try:
                if gate in _1Q:
                    sv = _apply_1q(sv, _1Q[gate], op[1], n)
                elif gate in ('rx','ry','rz'):
                    mat = {'rx':_Rx,'ry':_Ry,'rz':_Rz}[gate](op[2])
                    sv = _apply_1q(sv, mat, op[1], n)
                elif gate == 'cnot':
                    sv = _apply_2q(sv, _CNOT(), op[1], op[2], n)
                elif gate == 'cz':
                    sv = _apply_2q(sv, _CZ(), op[1], op[2], n)
                elif gate == 'cy':
                    sv = _apply_2q(sv, _CY(), op[1], op[2], n)
                elif gate == 'swap':
                    sv = _apply_2q(sv, _SWAP(), op[1], op[2], n)
                elif gate == 'toffoli':
                    sv = _apply_3q(sv, _TOFF(), op[1], op[2], op[3], n)
            except Exception as e:
                print(f"  [Debugger] ⚠️  Step {step} ({gate}) failed: {e}")

            # Save snapshot at every step
            snap = {
                "step":        step,
                "gate_label":  raw[i] if i < len(raw) else gate,
                "statevector": sv.copy(),
                "probs":       np.abs(sv)**2,
                "phases":      np.angle(sv),
                "labels":      labels,
            }

            # Hit a breakpoint → print it
            if step in self._breakpoints:
                self._snapshots[step] = snap
                self._print_breakpoint(snap)

        hit = len(self._snapshots)
        print(f"\n[Debugger] Finished. {hit}/{len(self._breakpoints)} "
              f"breakpoint(s) hit.")
        return self._snapshots

    # ── Inspect ───────────────────────────────────────────────────

    def get_snapshot(self, step_index: int) -> dict:
        """
        Retrieve the saved snapshot at a given step.
        Only works for steps that had a breakpoint.

        Returns the snapshot dict or None if not found.
        """
        snap = self._snapshots.get(step_index)
        if snap is None:
            print(f"[Debugger] No snapshot at step {step_index}. "
                  f"Available: {sorted(self._snapshots.keys())}")
        return snap

    def inspect(self, step_index: int):
        """Print the snapshot at a given step (call after run_debug)."""
        snap = self.get_snapshot(step_index)
        if snap:
            self._print_breakpoint(snap)

    # ── Pretty printer ────────────────────────────────────────────

    def _print_breakpoint(self, snap):
        n_states = len(snap["labels"])
        n_qubits = int(np.log2(n_states))

        print(f"\n  ┌─ BREAKPOINT — Step {snap['step']} "
              f"| Gate: {snap['gate_label']}")
        print(f"  │  {'State':<8} {'Prob':>8}  {'Amplitude':>22}  "
              f"{'Phase(rad)':>10}")
        print(f"  │  {'-'*55}")

        for label, amp, prob, phase in zip(
                snap["labels"], snap["statevector"],
                snap["probs"],  snap["phases"]):

            if prob < 1e-9:
                continue

            bar = '█' * int(prob * 16) + '░' * (16 - int(prob * 16))
            print(f"  │  {label:<8} {prob:>7.4f}  "
                  f"{amp.real:>+10.5f}{amp.imag:>+10.5f}j  "
                  f"{phase:>10.4f}")

        print(f"  └─ Bar chart:")
        for label, prob in zip(snap["labels"], snap["probs"]):
            if prob < 1e-9:
                continue
            bar = '█' * int(prob * 30) + '░' * (30 - int(prob * 30))
            print(f"     {label}  {bar}  {prob:.4f}")
        print()

    def plot_snapshot(self, step_index: int):
        """
        Plot the probability distribution at a breakpoint step
        as a bar chart using matplotlib.
        """
        import matplotlib.pyplot as plt

        snap = self.get_snapshot(step_index)
        if snap is None:
            return

        labels = snap["labels"]
        probs  = snap["probs"]
        phases = snap["phases"]

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

        # Probability bar chart
        colors = ['steelblue' if p > 1e-9 else 'lightgrey' for p in probs]
        ax1.bar(labels, probs, color=colors, edgecolor='navy', linewidth=0.5)
        ax1.set_title(f"Step {step_index} — {snap['gate_label']}\nProbabilities")
        ax1.set_ylabel("Probability")
        ax1.set_ylim(0, 1.1)
        ax1.tick_params(axis='x', rotation=45)
        ax1.grid(axis='y', alpha=0.3)

        # Phase bar chart
        active = [(l, ph) for l, p, ph in zip(labels, probs, phases) if p > 1e-9]
        if active:
            al, aph = zip(*active)
            colors2 = ['tomato' if p >= 0 else 'steelblue' for p in aph]
            ax2.bar(al, aph, color=colors2, edgecolor='darkred', linewidth=0.5)
        ax2.set_title(f"Step {step_index} — Phase Angles (radians)")
        ax2.set_ylabel("Phase (radians)")
        ax2.axhline(0, color='black', linewidth=0.8)
        ax2.tick_params(axis='x', rotation=45)
        ax2.grid(axis='y', alpha=0.3)

        plt.suptitle(f"PKTron Debugger — Snapshot at Step {step_index}",
                     fontweight='bold')
        plt.tight_layout()
        plt.savefig(f'pktron_debug_step{step_index}.png',
                    dpi=120, bbox_inches='tight')
        plt.show()
        print(f"[Debugger] Plot saved as pktron_debug_step{step_index}.png")
