{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0978387c-7ee0-4f00-b52b-b4b705ffbe72",
   "metadata": {},
   "source": [
    "# Lab 5: Simulating an N-qubit system\n",
    "\n",
    "## Introduction: Building on Single Qubit Systems\n",
    "\n",
    "Last week, you built a `SingleQubitSystem` class that simulated quantum states and operations on a single qubit. You implemented core methods to initialize quantum states, inspect them (measuring probabilities), and apply quantum gates like Hadamard and Pauli-Z. This week, we're taking a significant step forward: you'll generalize that single-qubit simulator to handle an **arbitrary number of qubits**.\n",
    "\n",
    "In the previous assignment, you were provided with substantial skeleton code that outlined the class structure, method signatures, and helper functions. This scaffolding allowed you to focus on understanding the quantum mechanics and implementing the core logic. **In this assignment, you will be expected to do everything yourself.** You'll implement all required methods, and write comprehensive tests.\n",
    "\n",
    "\n",
    "The jump from single-qubit to N-qubit systems requires a fundamental shift in how you think about quantum state representation.\n",
    "\n",
    "In the `SingleQubitSystem`, you could hardcode quantum gate matrices—a 2×2 matrix for single-qubit gates is simple and manageable. Now, with N qubits, gate matrices scale exponentially: an operation on an N-qubit system requires a 2^N × 2^N matrix. Rather than hardcoding these, you'll need to **programmatically construct** gate matrices and tensor products that generalize to any number of qubits.\n",
    "\n",
    "Beyond extending existing operations, in this lab you will also implement:\n",
    "\n",
    "- **Multi-qubit gates**: You'll implement the SWAP and the CNOT gates.\n",
    "- **Quantum oracles**: You'll implement two oracles—the **BernVaz oracle** and the **Archimedes oracle**.\n",
    "\n",
    "# Task 0: Swapping bits\n",
    "\n",
    "Before we dive into quantum mechanics, let's build intuition for how **SWAP** operations work at the bit level. Understanding classical bit swapping will help you reason about the quantum SWAP gate you'll implement later.\n",
    "\n",
    "We've provided you with a helper function that converts integers to binary strings in **big-endian**:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4261a6de-c296-4aff-ba2b-2dd8e6299094",
   "metadata": {},
   "outputs": [],
   "source": [
    "def int_to_bin_string(n: int, number_of_bits: int) -> str:\n",
    "    \"\"\"\n",
    "    Convert an integer to a binary string in big-endian format.\n",
    "\n",
    "    Args:\n",
    "        n: The integer to convert\n",
    "        number_of_bits: The length of the output binary string\n",
    "\n",
    "    Returns:\n",
    "        A binary string of length number_of_bits in big-endian format\n",
    "\n",
    "    \"\"\"\n",
    "    return format(n, f'0{number_of_bits}b')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec97f5bf-918a-4781-a0de-0f93f01603a1",
   "metadata": {},
   "source": [
    "\n",
    "A 4-bit system (called a **nibble**) can represent 16 different states (0 through 15). Your task is to print a truth table showing what happens when you swap the bits at **index 0** and **index 2**.\n",
    "\n",
    "The truth table should have three columns:\n",
    "\n",
    "1. **Original**: The original 4-bit string\n",
    "2. **After Swap**: The 4-bit string after swapping bits at positions 0 and 2\n",
    "3. **Decimal**: The decimal value of the resulting bit string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85cb044c-efbd-44ce-9f35-d0d08b010d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Format your table like this:\n",
    "original = after_swap = int_to_bin_string(3,4)\n",
    "print(f'{original}\\t{after_swap}')\n",
    "\n",
    "## TODO: Print the truth table after swapping bits at positions 0 and 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5e22504-9f84-4400-a02d-0b920b5db601",
   "metadata": {},
   "source": [
    "\n",
    "Now that you've seen how bit swapping works at the classical level, think about how this concept might extend to quantum systems.\n",
    "\n",
    "Consider these questions as you reflect:\n",
    "\n",
    "1. **State representation**: In your `SingleQubitSystem`, you represented a quantum state as a state vector—a list of amplitudes, one for each possible classical state. How might you use bit swapping to rearrange the elements of an N-qubit state vector?\n",
    "2. **Permutation and basis states**: When you swap bits in a classical state, you're essentially permuting which basis state corresponds to which computational outcome. If bit swapping changes the binary representation, how does that change *which* element of the state vector you're modifying?\n",
    "3. **Generalizing to multiple qubits**: In a 2-qubit system, there are 4 basis states ($|00\\rangle$, $|01\\rangle$, $|10\\rangle$, $|11\\rangle$). If you wanted to swap qubits 0 and 1, how would the indices of your state vector elements need to be rearranged? Can you see a pattern that would work for *any* pair of qubits in an N-qubit system?\n",
    "4. **Efficiency**: Rather than constructing a full $2^N \\times 2^N$ matrix, could you implement SWAP by directly rearranging elements of a state vector? What would be the computational advantage?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "843e73fe-2ba5-4d62-a93c-3440db997d16",
   "metadata": {},
   "source": [
    "\n",
    "## `NQubitSystem` technical specs\n",
    "\n",
    "`NQubitSystem` is a concrete implementation that inherits from the abstract `QubitSystem` base class. It generalizes single-qubit operations to support quantum systems with an arbitrary number of qubits.\n",
    "\n",
    "\n",
    "| Function | Behavior |\n",
    "| :-- | :-- |\n",
    "| `__init__(num_qubits: int)` | Initializes an N-qubit system with `num_qubits` qubits. The initial state should be the computational basis state $\\lvert 0 \\ldots 0 \\rangle$ (all qubits in state $\\lvert 0 \\rangle$). Raise a `ValueError` if `num_qubits` is less than 1. |\n",
    "| `set_value(value: list[float])` | Sets the system's quantum state. Accept a list of floats representing the state vector amplitudes in **big-endian** basis order. The length of the input must equal $2^{\\text{num_qubits}}$. Implementations should validate that the input represents a valid state (i.e., the sum of squared amplitudes equals 1.0, within numerical precision) and raise a `ValueError` if not. |\n",
    "| `get_value_braket() -> str` | Returns a bra-ket style string representing the current state. This is primarily for readability/debugging, you won't be penalized for stylistic differences or formatting. |\n",
    "| `get_value_vector() -> list[float]` | Returns the current state vector amplitudes in big-endian basis order. |\n",
    "| `apply_not(i: int) -> None` | Applies the NOT gate (Pauli-$X$) to qubit at index $i$, updating the system state. Raise an `IndexError` if $i$ is not a valid qubit index. |\n",
    "| `apply_h(i: int) -> None` | Applies the Hadamard gate ($H$) to qubit at index $i$, updating the system state. Raise an `IndexError` if $i$ is not a valid qubit index. |\n",
    "| `apply_z(i: int) -> None` | Applies the Pauli-$Z$ gate to qubit at index $i$, updating the system state. Raise an `IndexError` if $i$ is not a valid qubit index. |\n",
    "| `apply_cnot(control: int, target: int) -> None` | Applies a controlled-NOT gate with `control` as the control qubit and `target` as the target qubit, updating the system state. Raise an `IndexError` if either index is invalid, or if they are the same. |\n",
    "| `apply_swap(i: int, j: int) -> None` | Swaps qubits at indices $i$ and $j$, updating the system state. Raise an `IndexError` if either index is invalid, or if they are the same. |\n",
    "| `measure() -> str` | <p>Simulates a measurement of the state of the system and returns one of the possible values as a big-endian string of binary. For example, if the state is $\\lvert 000 \\rangle$, the result would always be `'000'`. If the state is $\\frac{1}{\\sqrt{2}} \\lvert 000 \\rangle + \\frac{1}{\\sqrt{2}} \\lvert 101 \\rangle$, measurement would return `'000'` or `'101'` with equal probability (50% each).<ul><li>Note: If a system is in a state of superposition before `measure()`, the act of measurement should collapse the superposition.</li><li>Note 2: The output should always have the same number of bits as the system does. For a 3-qubit system, outputs will be 3-character strings like `'000'`, `'101'`, etc.</li></ul></p> |\n",
    "\n",
    "\n",
    "## Using Numpy\n",
    "\n",
    "\n",
    "As with last week's lab, you are free to use `numpy` to handle mathematical operations. You may **not** use `qiskit` or **any** other library besides `numpy` as part of your solution. **Adding extra `import` statements _will_ crash the autograder.**\n",
    "\n",
    "You can perform vector-matrix multiplication with `numpy` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "998a01d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Define a state vector (e.g., |01⟩ for a 2-qubit system)\n",
    "state = np.array([0, 1, 0, 0])\n",
    "\n",
    "# Define a gate (e.g., a 2-qubit gate, X on the second qubit)\n",
    "gate = np.array([[0, 1, 0, 0],\n",
    "                 [1, 0, 0, 0],\n",
    "                 [0, 0, 0, 1],\n",
    "                 [0, 0, 1, 0]])\n",
    "\n",
    "# Apply the gate to the state using matrix-vector multiplication (Note: Numpy treats vectors as row-vectors)\n",
    "new_state = gate @ state\n",
    "print(new_state)\n",
    "# Output: [1 0 0 0]\n",
    "\n",
    "# Alternatively, using np.dot()\n",
    "new_state = np.dot(gate, state)\n",
    "print(new_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2a09c9e",
   "metadata": {},
   "source": [
    "To compute the tensor product (also called the Kronecker product):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6800e7d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Define two small matrices (e.g., single-qubit gates)\n",
    "I = np.array([[1, 0],\n",
    "              [0, 1]])\n",
    "\n",
    "X = np.array([[0, 1],\n",
    "              [1, 0]])\n",
    "\n",
    "# Tensor product\n",
    "result = np.kron(I, X)\n",
    "print(result)\n",
    "\n",
    "# Output:\n",
    "#[[0 1 0 0]\n",
    "# [1 0 0 0]\n",
    "# [0 0 0 1]\n",
    "# [0 0 1 0]]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfc4b39a",
   "metadata": {},
   "source": [
    "## Task 1: Copy over your `QubitSystem` implementation from last week's assignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2a1801b-2e76-482a-af4e-e0bee4222149",
   "metadata": {},
   "outputs": [],
   "source": [
    "from abc import ABC, abstractmethod\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "## TODO: Paste QubitSystem here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f25cef5a-f5d0-40a4-96bf-1ceb897c4a96",
   "metadata": {},
   "source": [
    "## Task 2: Adapt your test suite from `SingleQubitSystem` for `NQubitSystem`\n",
    "Just like last week, the autograder will expect the following tests to be present in your submission: `test_set_value`,`test_get_value_vector`,`test_apply_not`,`test_apply_h`,`test_apply_z`,`test_apply_cnot`,`test_apply_swap`,`test_measure`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5951bb0-3a43-46d8-b4cd-69680a9f7552",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: Implement your test suite here\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f48ab5-265a-4563-be7f-83ef183816fd",
   "metadata": {},
   "source": [
    "## Task 3: Implement `NQubitSystem`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4ba82b-eb64-4262-82bc-44257df5fed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: Implement NQubitSystem Here\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e45c6c-40df-496a-8662-c55a2be92fb7",
   "metadata": {},
   "source": [
    "## Oracles:\n",
    "\n",
    "\n",
    "In class we've discussed two quantum oracles: **BernVaz** and **Archimedes**. Notice how they share a common pattern:\n",
    "\n",
    "1. **State Preparation**: Both oracles expect the system to be in a specific configuration. BernVaz, for example, expects all qubits must be in a superposition state (e.g., created by applying Hadamard gates to all qubits).\n",
    "2. **Oracle Probing**: Once prepared, you apply the oracle itself, which transforms the superposed state according to the code(s) it encodes.\n",
    "3. **Post-Processing**: After the oracle runs, the system requires additional quantum operations to extract useful information. In the case of BernVaz, this means applying Hadamard gates to all qubits again to collapse the superposition in a meaningful way.\n",
    "\n",
    "Notice how both oracles follow this same high-level workflow, even though they encode different functions. This is precisely the kind of scenario where **object-oriented programming** shines.\n",
    "\n",
    "## Interfaces: Capturing Shared Behavior\n",
    "\n",
    "In OOP, an **interface** (called an abstract base class in python) defines a contract: \"Any object that implements this interface must support these operations.\" Interfaces let you write code that works with *any* object following that contract, without needing to know the specific details of each implementation.\n",
    "\n",
    "In our case, you might imagine defining an `Oracle` interface that specifies:\n",
    "\n",
    "- An oracle must be able to prepare the system\n",
    "- An oracle must be able to apply itself to a system\n",
    "- An oracle must be able to post-process the result\n",
    "\n",
    "Both the BernVaz and Archimedes oracles could implement this interface, each providing their own specific logic while adhering to the same structure. This way, future code could work with either oracle.\n",
    "\n",
    "This pattern is powerful because it lets you **abstract away the differences** (number of qubits, specific gate sequences) while **capturing the similarities** (the three-step workflow). As quantum algorithms grow more complex, this kind of abstraction will help you manage complexity and extend your code without rewriting it.\n",
    "\n",
    "# `Oracle` Abstract Base Class\n",
    "\n",
    "Implement `Oracle` as an abstract base class (an `ABC` from the Python library `abc`) that defines the interface for quantum oracles (this means `Oracle` should consist entirely of `@abstractmethod`s). It is up to concrete oracle implementations to inherit from this class and provide concrete implementations of the three abstract methods.\n",
    "\n",
    "\n",
    "| Function | Behavior |\n",
    "| :-- | :-- |\n",
    "| `__init__(codes: list[str])` | Initializes the oracle with a list of strings representing the 3-bit codes that define the oracle's behavior. Store these codes for use in the `probe()` method. |\n",
    "| `pre_probe(system: NQubitSystem) -> None` | Prepares the quantum system into the state that the oracle expects. This typically involves creating an equal superposition across all qubits. Raise an `IndexError` if `system.num_qubits` is not equal to 4. |\n",
    "| `probe(system: NQubitSystem) -> None` | Applies the oracle transformation to the system. The oracle uses the stored codes to determine how to transform each basis state. Raise an `IndexError` if `system.num_qubits` is not equal to 4. |\n",
    "| `post_probe(system: NQubitSystem) -> None` | Performs post-processing on the system after the oracle has been applied. This typically involves applying additional gates (e.g., Hadamard gates) to extract meaningful information from the superposed state. Raise an `IndexError` if `system.num_qubits` is not equal to 4. |\n",
    "\n",
    "## Task 4: Implement `Oracle` as an abstract base class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c476e7e8-873a-4d9f-9363-f5125c4a97c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: Oracle code below\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f1afcca-b10c-42f0-bd96-95d5f2f0b089",
   "metadata": {},
   "source": [
    "## Task 5: Implement the `BernVaz` oracle as a child of `Oracle`\n",
    "![bern vaz oracle](https://www.classes.cs.uchicago.edu/archive/2026/winter/22880-1/assigns/week5/figs/bernvazoraclealg.png)\n",
    "\n",
    "**Note**: We will only test your code with code lists containing a single 3 bit code (e.g., `['011']`). It is up to you to decide how to handle the case of lists containing multiple codes; you may choose to ignore that case or throw an `IndexError` exception."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac41c41-0c19-4c4c-8adb-00992fed1bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: BernVaz code below\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19630c53-61aa-4aaf-836f-f70f2d4b87f6",
   "metadata": {},
   "source": [
    "## Task 6: Implement the `Archimedes` oracle as a child of `Oracle`\n",
    "![archimedes oracle](https://www.classes.cs.uchicago.edu/archive/2026/winter/22880-1/assigns/week5/figs/archoraclealg.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3db6b5c-70f4-45c7-b28e-53bd386e9c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: Archimedes code below\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c9342d9-0f0b-423b-babc-431fadfcbd05",
   "metadata": {},
   "source": [
    "## Task 7: Test Your Oracles:\n",
    "\n",
    "Finally, write tests to verify that your `Oracle` abstract base class and concrete oracle implementations work correctly. Implement the following:\n",
    "\n",
    "\n",
    "| Test Name | Description |\n",
    "| :-- | :-- |\n",
    "| `test_oracle_is_abstract` | Verify that attempting to instantiate `Oracle` directly raises a `TypeError`. Since `Oracle` is an abstract base class, it should not be possible to create an instance of it without implementing all abstract methods. |\n",
    "| `test_bernvaz_is_oracle` | Verify that `BernVazOracle` is a subclass of `Oracle` and can be instantiated with a valid list of 3-bit codes. |\n",
    "| `test_archimedes_is_oracle` | Verify that `ArchimedesOracle` is a subclass of `Oracle` and can be instantiated with a valid list of 3-bit codes. |\n",
    "| `test_bernvaz` | Create a 4-qubit `NQubitSystem`, prepare it using `BernVazOracle.pre_probe()`, apply the oracle using `BernVazOracle.probe()`, perform post-processing with `BernVazOracle.post_probe()`, then verify that the final state matches your expected result. Test with multiple different code lists to ensure correctness. |\n",
    "| `test_archimedes` | Create a 4-qubit `NQubitSystem`, prepare it using `ArchimedesOracle.pre_probe()`, apply the oracle using `ArchimedesOracle.probe()`, perform post-processing with `ArchimedesOracle.post_probe()`, then verify that the final state matches your expected result. Test with multiple different code lists to ensure correctness. |\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716f2298-5700-41a4-b2f4-33b26823817e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO: Oracle Tests\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4877a4b6-ca8f-46d2-889a-4b3c03e12777",
   "metadata": {},
   "source": [
    "# Submission\n",
    "Congratulations on completing the lab!\n",
    "Make sure you:\n",
    "\n",
    "\n",
    "1.   Test all of your functions by calling them at least once.\n",
    "2.   Download your lab as a **python** `.py` script (NOT an `.ipynb` file):\n",
    "      \n",
    "      ```File -> Download -> Download .py```\n",
    "\n",
    "3.   Rename the downloaded file to `Lab5Answers.py`.\n",
    "4.   Upload the `Lab5Answers.py` file to Gradescope.\n",
    "5.   Ensure the autograder runs successfully."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
