
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Step-by-Step Simulation Engine
#
#  Runs your circuit ONE GATE AT A TIME.
#  After every single gate it records:
#    - the full statevector  (all complex amplitudes)
#    - the probabilities     (chance of measuring each basis state)
#
#  Perfect for learning what each gate actually does!
# ─────────────────────────────────────────────────────────────────

# ── Gate matrices ─────────────────────────────────────────────────
_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

# ── Low-level gate application (same logic as PKTron core) ────────

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: turn gate_history strings into structured ops ─────────

def _parse_history(gate_history):
    """
    Convert PKTron gate_history strings into (gate, *args) tuples.
    Handles formats like:
      "H on qubit 0"
      "CNOT control 0 target 1"
      "Rx(1.57) on qubit 0"
      "Toffoli control1 0 control2 1 target 2"
      "SWAP qubits 0 and 1"
    """
    import re
    ops = []
    for entry in gate_history:
        e = entry.strip()
        el = e.lower()

        # Rx / Ry / Rz  — e.g. "Rx(1.57) on qubit 0"
        m = re.match(r'r([xyz])\(([^)]+)\) on qubit (\d+)', el)
        if m:
            gate = 'r' + m.group(1)
            theta = float(m.group(2))
            q = int(m.group(3))
            ops.append((gate, q, theta))
            continue

        # H/X/Y/Z/S/T and their daggers — "H on qubit 0"
        m = re.match(r'([a-z†]+) on qubit (\d+)', el)
        if m:
            ops.append((m.group(1).replace('†','dg'), int(m.group(2))))
            continue

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

        # CZ / CY
        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

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

        # Toffoli — "toffoli control1 0 control2 1 target 2"
        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

        # S† / T†
        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 function ─────────────────────────────────────────────────

def run_step_by_step(circuit) -> list:
    """
    Execute a PKTron circuit gate-by-gate.

    Returns a list of dicts, one per gate:
      {
        "step"       : step number (1-based),
        "gate_label" : human-readable gate name from gate_history,
        "statevector": numpy complex array,
        "probs"      : numpy float array (probabilities),
        "labels"     : basis state labels e.g. ["|00>", "|01>", ...]
      }

    Example:
        from pktron import generate_circuit
        from pktron.step_simulator import run_step_by_step, print_steps

        qc = generate_circuit("bell state")
        steps = run_step_by_step(qc)
        print_steps(steps)
    """
    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)]

    # Start from |00...0>
    sv = np.zeros(2**n, dtype=complex)
    sv[0] = 1.0

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

    results = []

    # Record initial state as step 0
    results.append({
        "step": 0,
        "gate_label": "Initial state |" + "0"*n + ">",
        "statevector": sv.copy(),
        "probs": np.abs(sv)**2,
        "labels": labels,
    })

    for i, op in enumerate(ops):
        gate = op[0].lower()

        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"  [StepSim] Skipped step {i+1} ({gate}): {e}")

        results.append({
            "step": i + 1,
            "gate_label": raw[i] if i < len(raw) else gate,
            "statevector": sv.copy(),
            "probs": np.abs(sv)**2,
            "labels": labels,
        })

    return results


def print_steps(steps, show_zero: bool = False):
    """
    Pretty-print every step from run_step_by_step().

    Args:
        steps:     list returned by run_step_by_step()
        show_zero: if False, hide basis states with 0 probability
    """
    total = len(steps) - 1   # step 0 is initial state
    print("=" * 60)
    print("  PKTron Step-by-Step Simulation")
    print(f"  {total} gate(s) | {len(steps[0]['labels'])} basis states")
    print("=" * 60)

    for s in steps:
        print(f"\n  ── Step {s['step']} : {s['gate_label']}")
        print(f"  {'State':<8}  {'Probability':>11}  {'Bar':}")
        print(f"  {'-'*45}")
        for label, prob in zip(s['labels'], s['probs']):
            if not show_zero and prob < 1e-9:
                continue
            bar = '█' * int(prob * 24) + '░' * (24 - int(prob * 24))
            print(f"  {label:<8}  {prob:>10.4f}   {bar}")

    print("\n" + "=" * 60)


def plot_steps(steps):
    """
    Plot probability bar charts for every step using matplotlib.
    Each gate gets its own subplot.
    """
    import matplotlib.pyplot as plt

    n_steps = len(steps)
    cols = min(n_steps, 3)
    rows = (n_steps + cols - 1) // cols

    fig, axes = plt.subplots(rows, cols,
                             figsize=(5 * cols, 3.5 * rows))
    axes = np.array(axes).reshape(-1)   # flatten for easy indexing

    for i, s in enumerate(steps):
        ax = axes[i]
        labels = s['labels']
        probs  = s['probs']
        colors = ['steelblue' if p > 1e-9 else 'lightgrey' for p in probs]
        ax.bar(labels, probs, color=colors, edgecolor='navy', linewidth=0.5)
        ax.set_title(f"Step {s['step']}\n{s['gate_label']}", fontsize=8)
        ax.set_ylim(0, 1.05)
        ax.set_ylabel('Probability', fontsize=7)
        ax.tick_params(axis='x', rotation=45, labelsize=7)
        ax.grid(axis='y', alpha=0.3)

    # Hide unused subplots
    for j in range(n_steps, len(axes)):
        axes[j].set_visible(False)

    plt.suptitle('PKTron — Step-by-Step Simulation', fontsize=11,
                 fontweight='bold', y=1.01)
    plt.tight_layout()
    plt.savefig('pktron_steps.png', dpi=120, bbox_inches='tight')
    plt.show()
    print("[StepSim] Plot saved as pktron_steps.png")
