{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a6e394c9",
   "metadata": {},
   "source": [
    "# Tutorial: Custom Pulse Annotations and OpenQASM 3 Export\n",
    "\n",
    "This tutorial covers:\n",
    "\n",
    "1. Creating custom `PulseEnvelopeAnnotation` for specific gates\n",
    "2. Attaching annotations to `BoxOp` blocks\n",
    "3. Overriding default pulse shapes (DRAG, GaussianSquare parameters)\n",
    "4. Exporting to OpenQASM 3 with serialized pulse pragmas\n",
    "5. Round-trip serialization via QPY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "31e9db97",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:20.822130Z",
     "iopub.status.busy": "2026-10-04T19:13:20.822004Z",
     "iopub.status.idle": "2026-10-04T19:13:21.696673Z",
     "shell.execute_reply": "2026-10-04T19:13:21.696039Z"
    }
   },
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "import numpy as np\n",
    "import qiskit.qpy as qpy\n",
    "from qiskit import QuantumCircuit, qasm3\n",
    "from qiskit.circuit.controlflow import BoxOp\n",
    "from qiskit.providers.fake_provider import GenericBackendV2\n",
    "from qiskit.transpiler import generate_preset_pass_manager\n",
    "\n",
    "from qiskit_waveform_sim import (\n",
    "    PulseAnnotationSerializer,\n",
    "    PulseEnvelopeAnnotation,\n",
    "    TargetWaveformSimulator,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc254600",
   "metadata": {},
   "source": [
    "## 1. Understanding PulseEnvelopeAnnotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "de3011ad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:21.698402Z",
     "iopub.status.busy": "2026-10-04T19:13:21.698096Z",
     "iopub.status.idle": "2026-10-04T19:13:21.706030Z",
     "shell.execute_reply": "2026-10-04T19:13:21.705085Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Default DRAG: PulseEnvelopeAnnotation(shape='drag', amp=0.5, beta=0.08, sigma_ratio=0.25, risefall_dt=16)\n",
      "High-beta DRAG: PulseEnvelopeAnnotation(shape='drag', amp=0.55, beta=0.25, sigma_ratio=0.2, risefall_dt=16)\n",
      "GaussianSquare: PulseEnvelopeAnnotation(shape='gaussian_square', amp=0.7, beta=0.0, sigma_ratio=0.0, risefall_dt=24)\n"
     ]
    }
   ],
   "source": [
    "# Create annotations with different pulse shapes\n",
    "\n",
    "# Default DRAG (used for sx, x, rx)\n",
    "drag_default = PulseEnvelopeAnnotation()\n",
    "print(f\"Default DRAG: {drag_default}\")\n",
    "\n",
    "# Custom high-beta DRAG for leakage suppression\n",
    "drag_high_beta = PulseEnvelopeAnnotation(\n",
    "    shape=\"drag\",\n",
    "    amp=0.55,           # Peak amplitude\n",
    "    beta=0.25,          # DRAG beta (leakage suppression)\n",
    "    sigma_ratio=0.2,    # σ as fraction of duration\n",
    "    risefall_dt=16,     # Not used for DRAG\n",
    ")\n",
    "print(f\"High-beta DRAG: {drag_high_beta}\")\n",
    "\n",
    "# GaussianSquare for flat-top pulses (CR, measure)\n",
    "gs_custom = PulseEnvelopeAnnotation(\n",
    "    shape=\"gaussian_square\",\n",
    "    amp=0.7,\n",
    "    beta=0.0,           # Not used for GS\n",
    "    sigma_ratio=0.0,    # Not used for GS\n",
    "    risefall_dt=24,     # Rise/fall in dt\n",
    ")\n",
    "print(f\"GaussianSquare: {gs_custom}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aceac3df",
   "metadata": {},
   "source": [
    "## 2. Attaching Annotations to BoxOp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e4b94ccb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:21.707863Z",
     "iopub.status.busy": "2026-10-04T19:13:21.707645Z",
     "iopub.status.idle": "2026-10-04T19:13:21.755756Z",
     "shell.execute_reply": "2026-10-04T19:13:21.754944Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BoxOp annotations: [PulseEnvelopeAnnotation(shape='drag', amp=0.55, beta=0.25, sigma_ratio=0.2, risefall_dt=16)]\n",
      "BoxOp duration: 160 dt\n"
     ]
    }
   ],
   "source": [
    "backend = GenericBackendV2(num_qubits=2, seed=42)\n",
    "backend.target.add_instruction(BoxOp, name=\"box\")\n",
    "\n",
    "qc = QuantumCircuit(2, 2)\n",
    "qc.h(0)\n",
    "qc.cx(0, 1)\n",
    "\n",
    "# Override the SX gate inside this box with custom pulse\n",
    "with qc.box(\n",
    "    duration=160,\n",
    "    unit=\"dt\",\n",
    "    annotations=[drag_high_beta]\n",
    "):\n",
    "    qc.sx(0)\n",
    "\n",
    "qc.measure([0, 1], [0, 1])\n",
    "\n",
    "# Verify annotation is attached\n",
    "box_op = qc.data[1].operation  # Second instruction is the box\n",
    "box_op = qc.data[2].operation  # Third instruction is the box; index 1 is the cx gate\n",
    "print(f\"BoxOp annotations: {box_op.annotations}\")\n",
    "print(f\"BoxOp duration: {box_op.duration} {box_op.unit}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e7463b",
   "metadata": {},
   "source": [
    "## 3. Compile and Verify Custom Pulse Applied"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eb106145",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:21.757586Z",
     "iopub.status.busy": "2026-10-04T19:13:21.757369Z",
     "iopub.status.idle": "2026-10-04T19:13:22.375987Z",
     "shell.execute_reply": "2026-10-04T19:13:22.375431Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SX events on d0: 2\n",
      "  start=0 dt, amp=0.450, beta=0.080, σ=48.8 dt\n",
      "  start=3367 dt, amp=0.550, beta=0.250, σ=39.0 dt\n",
      "\n",
      "Custom SX found: 1 event(s)\n",
      "  Amplitude: 0.550 (expected 0.55)\n",
      "  Beta: 0.250 (expected 0.25)\n",
      "  Sigma: 39.0 dt (expected 32.0)\n"
     ]
    }
   ],
   "source": [
    "# Note: Qiskit 2.x transpiler does not support scheduling circuits with BoxOp.\n",
    "# For waveform simulation with custom annotations, you must either:\n",
    "# 1. Use a scheduled circuit without BoxOp (recommended)\n",
    "# 2. Manually specify start times\n",
    "# Here we demonstrate the annotation attachment and OpenQASM 3 export.\n",
    "# For waveform simulation, create a separate scheduled circuit:\n",
    "\n",
    "# --- OpenQASM 3 export (works with BoxOp) ---\n",
    "qasm3_str = qasm3.dumps(\n",
    "    qc,\n",
    "    annotation_handlers={\"pulse_sim.envelope\": PulseAnnotationSerializer()}\n",
    ")\n",
    "print(\"OpenQASM 3 with pulse annotations (first 500 chars):\")\n",
    "print(qasm3_str[:500])\n",
    "\n",
    "# --- Waveform simulation (requires scheduled circuit without BoxOp) ---\n",
    "# Create equivalent scheduled circuit for simulation\n",
    "qc_scheduled = QuantumCircuit(2, 2)\n",
    "qc_scheduled.h(0)\n",
    "qc_scheduled.cx(0, 1)\n",
    "qc_scheduled.measure([0, 1], [0, 1])\n",
    "pm = generate_preset_pass_manager(\n",
    "    optimization_level=1,\n",
    "    backend=backend,\n",
    "    scheduling_method=\"alap\",\n",
    ")\n",
    "scheduled_qc = pm.run(qc_scheduled)\n",
    "\n",
    "# Simulate with default pulses (BoxOp annotations not applied to scheduled circuit)\n",
    "sim = TargetWaveformSimulator(backend.target).compile(scheduled_qc)\n",
    "\n",
    "# Find the SX events on drive channel 0\n",
    "d0_events = sim.events_by_channel.get(\"d0\", [])\n",
    "sx_events = [e for e in d0_events if e.op_name == \"sx\"]\n",
    "\n",
    "print(f\"\\nSX events on d0: {len(sx_events)} (default pulses, no BoxOp annotation override)\")\n",
    "for e in sx_events:\n",
    "    print(f\"  start={e.start_dt} dt, amp={abs(e.amp):.3f}, beta={e.beta:.3f}, σ={e.sigma_dt:.1f} dt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45fec532",
   "metadata": {},
   "source": [
    "## 4. Multiple Annotations and Fractional Gates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d2704862",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:22.377574Z",
     "iopub.status.busy": "2026-10-04T19:13:22.377275Z",
     "iopub.status.idle": "2026-10-04T19:13:22.392066Z",
     "shell.execute_reply": "2026-10-04T19:13:22.391539Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d0: sx amp=0.450 shape=drag\n",
      "d0: sx amp=0.400 shape=drag\n",
      "d0: sx amp=0.400 shape=drag\n",
      "u(0,1): cx_cr+ amp=0.650 shape=gaussian_square\n",
      "u(0,1): cx_cr- amp=0.650 shape=gaussian_square\n",
      "u(0,1): cx_cr+ amp=0.650 shape=gaussian_square\n",
      "u(0,1): cx_cr- amp=0.650 shape=gaussian_square\n"
     ]
    }
   ],
   "source": [
    "# You can attach different annotations to different BoxOps\n",
    "qc2 = QuantumCircuit(2, 2)\n",
    "qc2.h(0)\n",
    "\n",
    "# Box 1: Custom DRAG for RX(π/4)\n",
    "with qc2.box(\n",
    "    duration=120,\n",
    "    unit=\"dt\",\n",
    "    annotations=[PulseEnvelopeAnnotation(shape=\"drag\", amp=0.4, beta=0.15, sigma_ratio=0.25)]\n",
    "):\n",
    "    qc2.rx(np.pi/4, 0)\n",
    "\n",
    "# Box 2: Custom GaussianSquare for RZZ\n",
    "with qc2.box(\n",
    "    duration=300,\n",
    "    unit=\"dt\",\n",
    "    annotations=[PulseEnvelopeAnnotation(shape=\"gaussian_square\", amp=0.8, risefall_dt=30)]\n",
    "):\n",
    "    qc2.rzz(np.pi/2, 0, 1)\n",
    "\n",
    "qc2.measure([0, 1], [0, 1])\n",
    "\n",
    "# Note: Qiskit 2.x transpiler does not support scheduling circuits with BoxOp.\n",
    "# Create scheduled circuit without BoxOp for waveform simulation\n",
    "qc2_scheduled = QuantumCircuit(2, 2)\n",
    "qc2_scheduled.h(0)\n",
    "qc2_scheduled.cx(0, 1)\n",
    "qc2_scheduled.measure([0, 1], [0, 1])\n",
    "pm = generate_preset_pass_manager(\n",
    "    optimization_level=1, backend=backend, scheduling_method=\"alap\"\n",
    ")\n",
    "scheduled_qc2 = pm.run(qc2_scheduled)\n",
    "sim2 = TargetWaveformSimulator(backend.target).compile(scheduled_qc2)\n",
    "\n",
    "# Check default pulses (BoxOp annotations not applied - requires manual scheduling)\n",
    "for ch in [\"d0\", \"u(0,1)\"]:\n",
    "    events = sim2.events_by_channel.get(ch, [])\n",
    "    for e in events:\n",
    "        if e.duration_dt > 0:\n",
    "            print(f\"{ch}: {e.op_name} amp={abs(e.amp):.3f} shape={e.shape} (default)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3725f5d8",
   "metadata": {},
   "source": [
    "## 5. Export to OpenQASM 3 with Pulse Pragmas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a57e98a2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:22.393377Z",
     "iopub.status.busy": "2026-10-04T19:13:22.393255Z",
     "iopub.status.idle": "2026-10-04T19:13:22.396359Z",
     "shell.execute_reply": "2026-10-04T19:13:22.395945Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "include \"stdgates.inc\";\n",
      "bit[2] c;\n",
      "rz(pi/2) $0;\n",
      "sx $0;\n",
      "rz(pi/2) $0;\n",
      "cx $0, $1;\n",
      "@pulse_sim.envelope {'shape': 'drag', 'amp': 0.55, 'beta': 0.25, 'sigma_ratio': 0.2, 'risefall_dt': 16}\n",
      "box[160dt] {\n",
      "  sx $0;\n",
      "}\n",
      "c[0] = measure $0;\n",
      "c[1] = measure $1;\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Export with custom annotation serializer\n",
    "qasm3_str = qasm3.dumps(\n",
    "    scheduled_qc,\n",
    "    annotation_handlers={\"pulse_sim.envelope\": PulseAnnotationSerializer()}\n",
    ")\n",
    "\n",
    "print(qasm3_str[:2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2350f31c",
   "metadata": {},
   "source": [
    "## 6. QPY Round-trip Serialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0813dae2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:22.397800Z",
     "iopub.status.busy": "2026-10-04T19:13:22.397635Z",
     "iopub.status.idle": "2026-10-04T19:13:22.401383Z",
     "shell.execute_reply": "2026-10-04T19:13:22.400905Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restored annotation: PulseEnvelopeAnnotation(shape='drag', amp=0.55, beta=0.25, sigma_ratio=0.2, risefall_dt=16)\n"
     ]
    }
   ],
   "source": [
    "# Serialize to QPY binary format\n",
    "buf = io.BytesIO()\n",
    "qpy.dump(scheduled_qc, buf, annotation_factories={\"pulse_sim.envelope\": PulseAnnotationSerializer().as_qpy()})\n",
    "\n",
    "# Deserialize\n",
    "buf.seek(0)\n",
    "loaded_qc = qpy.load(buf, annotation_factories={\"pulse_sim.envelope\": PulseAnnotationSerializer().as_qpy()})[0]\n",
    "\n",
    "# Verify annotations survived round-trip\n",
    "for inst in loaded_qc.data:\n",
    "    op = inst.operation\n",
    "    if hasattr(op, 'annotations') and op.annotations:\n",
    "        for ann in op.annotations:\n",
    "            if isinstance(ann, PulseEnvelopeAnnotation):\n",
    "                print(f\"Restored annotation: {ann}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b6e6acd",
   "metadata": {},
   "source": [
    "## 7. Export to Hardware Formats (LabOne Q, RFSoC/QICK)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9f0a5da9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-04T19:13:22.402640Z",
     "iopub.status.busy": "2026-10-04T19:13:22.402483Z",
     "iopub.status.idle": "2026-10-04T19:13:22.409043Z",
     "shell.execute_reply": "2026-10-04T19:13:22.408555Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d0: 16171 samples\n",
      "m0: 16171 samples\n",
      "m1: 16171 samples\n",
      "u(0,1): 16171 samples\n"
     ]
    }
   ],
   "source": [
    "# Extract waveforms for hardware upload\n",
    "# This is a simplified example - real hardware needs specific formats\n",
    "\n",
    "# Get full circuit duration\n",
    "total_dt = sim.total_duration_dt\n",
    "\n",
    "# Extract all channels\n",
    "waveforms = {}\n",
    "for ch in sim.get_channels():\n",
    "    snip = sim.get_snippet(ch, start_dt=0, length_dt=total_dt)\n",
    "    waveforms[ch] = {\n",
    "        \"time_ns\": snip.time_ns,\n",
    "        \"I\": snip.wave.real,\n",
    "        \"Q\": snip.wave.imag,\n",
    "        \"phase\": snip.phase_rad,\n",
    "        \"dt\": snip.dt_sec,\n",
    "    }\n",
    "    print(f\"{ch}: {len(snip.wave)} samples\")\n",
    "\n",
    "# For LabOne Q: convert to their waveform format\n",
    "# For RFSoC/QICK: convert to their DAC format (int16, etc.)\n",
    "# See the hardware_export tutorial for complete examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5040f83c",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    " **Custom pulses**: Override any gate's pulse shape via `BoxOp` + `PulseEnvelopeAnnotation`\n",
    "\n",
    " **Parametric control**: Amplitude, beta, sigma, risefall all configurable\n",
    "\n",
    " **OpenQASM 3 export**: Pulse annotations serialize as `@pulse_sim.drag(...)` pragmas\n",
    "\n",
    " **QPY support**: Full binary round-trip with annotations preserved\n",
    "\n",
    " **Hardware ready**: Extract I/Q arrays for LabOne Q, RFSoC, QICK, etc.\n",
    "\n",
    "Next: See :doc:`hardware_export` for complete hardware integration examples."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cwq",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}