diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index dc31eda..e1c3284 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -7,16 +7,16 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Check out code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - name: Install dependencies
       run: |
         pip install -U pip
         pip install -r requirements.txt
-        pip install "black>=21.5b0" "coverage>=5.5" "pytest>=2.1.0"
+        pip install black coverage pytest
     - name: Format code with Black
       run: python3 -m black $(find . -name '*.py')
     - name: Run tests with coverage
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 1892607..e02d362 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -21,9 +21,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - name: Install dependencies
@@ -33,7 +33,7 @@ jobs:
     - name: Build package
       run: python -m build
     - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      uses: pypa/gh-action-pypi-publish@release/v1
       with:
         user: __token__
         password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.gitignore b/.gitignore
index b8175c8..7dbdfee 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,7 +1,3 @@
-# Development
-Untitled.*
-untitled.*
-
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -80,9 +76,7 @@ docs/_build/
 target/
 
 # Jupyter Notebook
-.ipynb_checkpoints/
-.jupyter/
-.virtual_documents/
+.ipynb_checkpoints
 
 # IPython
 profile_default/
@@ -100,6 +94,12 @@ ipython_config.py
 #   install all needed dependencies.
 #Pipfile.lock
 
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
 # poetry
 #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
 #   This is especially recommended for binary packages to ensure reproducibility, and is more
@@ -112,8 +112,10 @@ ipython_config.py
 #pdm.lock
 #   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
 #   in version control.
-#   https://pdm.fming.dev/#use-with-ide
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
 .pdm.toml
+.pdm-python
+.pdm-build/
 
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
 __pypackages__/
@@ -164,3 +166,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# PyPI configuration file
+.pypirc
\ No newline at end of file
diff --git a/README.md b/README.md
index 8608b43..546ae4f 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # Qubit Simulator
 
-Qubit Simulator is a simple and lightweight library that provides a quantum simulator for simulating qubits and quantum gates. It supports basic quantum operations and gates such as Hadamard, π/8, Controlled-Not, and generic unitary transformations.
+Qubit Simulator is a simple and lightweight library that provides a quantum statevector simulator for simulating qubits and quantum gates. It supports basic quantum operations and gates using NumPy.
 
 ## Installation
 
@@ -37,7 +37,7 @@ simulator.cx(0, 2)  # Controlled-Not gate
 Define and apply custom gates using angles:
 
 ```python
-simulator.u(2, 0.3, 0.4, 0.5)  # Generic gate
+simulator.u(0.3, 0.4, 0.5, 2)  # Generic gate
 ```
 
 ### Measurements
@@ -52,32 +52,6 @@ print(simulator.run(shots=100))
 {'000': 46, '001': 4, '100': 4, '101': 46}
 ```
 
-### Circuit Representation
-
-Get a string representation of the circuit:
-
-```python
-print(simulator)
-```
-
-```plaintext
------------------------------------
-| H |   | @ |                     |
-|   | T |   |                     |
-|   |   | X | U(0.30, 0.40, 0.50) |
------------------------------------
-```
-
-### Wavefunction Plot
-
-Show the amplitude and phase of all quantum states:
-
-```python
-simulator.plot_wavefunction()
-```
-
-![Wavefunction Scatter Plot](https://github.com/splch/qubit-simulator/assets/25377399/de3242ef-9c14-44be-b49b-656e9727c618)
-
 ## Testing
 
 Tests are included in the package to verify its functionality and provide more advanced examples:
diff --git a/qubit_simulator/__init__.py b/qubit_simulator/__init__.py
index 5aed0b7..e10c2ba 100644
--- a/qubit_simulator/__init__.py
+++ b/qubit_simulator/__init__.py
@@ -1,2 +1,4 @@
 from .simulator import QubitSimulator
 from .gates import Gates
+
+__all__ = ["QubitSimulator", "Gates"]
diff --git a/qubit_simulator/gates.py b/qubit_simulator/gates.py
index c2b9b48..2850295 100644
--- a/qubit_simulator/gates.py
+++ b/qubit_simulator/gates.py
@@ -1,98 +1,69 @@
 import numpy as np
-from typing import Union
 
 
 class Gates:
-    """
-    A class that represents the quantum gates.
-    """
+    """Minimal collection of common gates and helper methods."""
 
-    # Hadamard (H) gate
-    H: np.ndarray = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
-    # π/8 (T) gate
-    T: np.ndarray = np.array([[1, 0], [0, np.exp(1j * np.pi / 4)]])
-    # Pauli-X (NOT) gate
-    X: np.ndarray = np.array([[0, 1], [1, 0]])
+    # Single-qubit gates (2x2)
+    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)
+    H = np.array([[1, 1], [1, -1]], dtype=complex) / np.sqrt(2)
+    S = np.diag([1, 1j]).astype(complex)
+    T = np.diag([1, np.exp(1j * np.pi / 4)]).astype(complex)
 
     @staticmethod
-    def U(theta: float, phi: float, lambda_: float) -> np.ndarray:
-        """
-        Generic (U) gate.
-
-        :param theta: Angle theta.
-        :param phi: Angle phi.
-        :param lambda_: Angle lambda.
-        :return: Unitary matrix representing the U gate.
-        """
+    def U(theta: float, phi: float, lam: float) -> np.ndarray:
         return np.array(
             [
-                [np.cos(theta), -np.exp(1j * lambda_) * np.sin(theta)],
+                [np.cos(theta / 2), -np.exp(1j * lam) * np.sin(theta / 2)],
                 [
-                    np.exp(1j * phi) * np.sin(theta),
-                    np.exp(1j * (phi + lambda_)) * np.cos(theta),
+                    np.exp(1j * phi) * np.sin(theta / 2),
+                    np.exp(1j * (phi + lam)) * np.cos(theta / 2),
                 ],
-            ]
+            ],
+            dtype=complex,
         )
 
+    # Two-qubit gates (4x4)
     @staticmethod
-    def create_controlled_gate(
-        gate: np.ndarray, control_qubit: int, target_qubit: int, num_qubits: int
-    ) -> np.ndarray:
-        """
-        Creates a controlled gate.
+    def SWAP_matrix() -> np.ndarray:
+        return np.array(
+            [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex
+        )
 
-        :param gate: Matrix representing the gate.
-        :param control_qubit: Index of the control qubit.
-        :param target_qubit: Index of the target qubit.
-        :param num_qubits: Total number of qubits.
-        :return: Matrix representing the controlled gate.
-        """
-        controlled_gate = np.eye(2**num_qubits, dtype=complex)
-        for basis in range(2**num_qubits):
-            basis_binary = format(basis, f"0{num_qubits}b")
-            if basis_binary[control_qubit] == "1":
-                target_state = int(
-                    basis_binary[:target_qubit]
-                    + str(1 - int(basis_binary[target_qubit]))
-                    + basis_binary[target_qubit + 1 :],
-                    2,
-                )
-                controlled_gate[basis, basis] = gate[
-                    int(basis_binary[target_qubit]), int(basis_binary[target_qubit])
-                ]
-                controlled_gate[basis, target_state] = gate[
-                    int(basis_binary[target_qubit]), 1 - int(basis_binary[target_qubit])
-                ]
-                controlled_gate[target_state, basis] = gate[
-                    1 - int(basis_binary[target_qubit]), int(basis_binary[target_qubit])
-                ]
-                controlled_gate[target_state, target_state] = gate[
-                    1 - int(basis_binary[target_qubit]),
-                    1 - int(basis_binary[target_qubit]),
-                ]
-        return controlled_gate
+    @staticmethod
+    def iSWAP_matrix() -> np.ndarray:
+        return np.array(
+            [[1, 0, 0, 0], [0, 0, 1j, 0], [0, 1j, 0, 0], [0, 0, 0, 1]], dtype=complex
+        )
 
+    # Three-qubit gates (8x8)
     @staticmethod
-    def create_inverse_gate(gate: np.ndarray) -> np.ndarray:
-        """
-        Creates an inverse gate.
+    def Toffoli_matrix() -> np.ndarray:
+        # Flip the 3rd qubit if first two are |1>
+        m = np.eye(8, dtype=complex)
+        m[[6, 7], [6, 7]] = 0
+        m[6, 7] = 1
+        m[7, 6] = 1
+        return m
 
-        :param gate: Matrix representing the gate.
-        :return: Matrix representing the inverse gate.
-        """
-        return np.conjugate(gate.T)
+    @staticmethod
+    def Fredkin_matrix() -> np.ndarray:
+        # Swap the last two qubits if the first is |1>
+        m = np.eye(8, dtype=complex)
+        m[[5, 6], [5, 6]] = 0
+        m[5, 6] = 1
+        m[6, 5] = 1
+        return m
 
     @staticmethod
-    def _validate_gate(gate: np.ndarray):
-        """
-        Validates the gate.
+    def inverse_gate(U: np.ndarray) -> np.ndarray:
+        return U.conjugate().T
 
-        :param gate: Matrix representing the gate.
-        :raises ValueError: If the gate is invalid.
-        """
-        if not np.allclose(
-            gate @ Gates.create_inverse_gate(gate), np.eye(gate.shape[0])
-        ):
-            raise ValueError(
-                "The gate must be unitary. Its conjugate transpose must be equal to its inverse."
-            )
+    @staticmethod
+    def controlled_gate(U: np.ndarray) -> np.ndarray:
+        c = np.zeros((4, 4), dtype=complex)
+        c[:2, :2] = np.eye(2)
+        c[2:, 2:] = U
+        return c
diff --git a/qubit_simulator/simulator.py b/qubit_simulator/simulator.py
index 64de396..83e29a0 100644
--- a/qubit_simulator/simulator.py
+++ b/qubit_simulator/simulator.py
@@ -1,296 +1,89 @@
 import numpy as np
-import matplotlib.pyplot as plt
-from collections import Counter
-from sys import getsizeof
-from typing import Optional, List, Tuple, Dict
 from .gates import Gates
 
 
 class QubitSimulator:
     """
-    A class that represents a qubit simulator.
+    Simple statevector simulator using tensor operations to apply
+    any k-qubit gate by appropriately reshaping both the gate and state.
     """
 
     def __init__(self, num_qubits: int):
-        """
-        Initialize the simulator with given number of qubits.
-
-        :param num_qubits: Number of qubits.
-        :raises ValueError: If the number of qubits is negative.
-        """
-        if num_qubits < 0:
-            raise ValueError("Number of qubits must be non-negative.")
-
-        self.num_qubits = num_qubits
-        self.state_vector = np.zeros(2**num_qubits, dtype=complex)
-        self.state_vector[0] = 1
-        self.circuit: List[Tuple[str, int, Optional[int]]] = []
-
-    def _validate_qubit_index(
-        self, target_qubit: int, control_qubit: Optional[int] = None
-    ):
-        """
-        Validates the qubit indices.
-
-        :param target_qubit: Index of the target qubit to validate.
-        :param control_qubit: Index of the control qubit to validate.
-        :raises IndexError: If the qubit index is out of range.
-        """
-        if target_qubit < 0 or target_qubit >= self.num_qubits:
-            raise IndexError(f"Target qubit index {target_qubit} out of range.")
-        if control_qubit is not None and (
-            control_qubit < 0 or control_qubit >= self.num_qubits
-        ):
-            raise IndexError(f"Control qubit index {control_qubit} out of range.")
-
-    def _get_gate_name(
-        self, theta: float, phi: float, lambda_: float, inverse: bool
-    ) -> str:
-        """
-        Constructs the name for a U gate or its inverse.
-
-        :param theta: Angle theta.
-        :param phi: Angle phi.
-        :param lambda_: Angle lambda.
-        :param inverse: Whether the gate is an inverse.
-        :return: String representing the gate name.
-        """
-        return f"U{'†' if inverse else ''}({theta:.2f}, {phi:.2f}, {lambda_:.2f})"
-
-    def _apply_gate(
-        self,
-        gate_name: str,
-        gate: np.ndarray,
-        target_qubit: int,
-        control_qubit: Optional[int] = None,
-    ):
-        """
-        Applies the given gate to the target qubit.
-
-        :param gate_name: Name of the gate.
-        :param gate: Matrix representing the gate.
-        :param target_qubit: Index of the target qubit.
-        :param control_qubit: Index of the control qubit (if controlled gate).
-        """
-        # Validate the target and control qubit indices
-        self._validate_qubit_index(target_qubit, control_qubit)
-        # Validate the gate
-        Gates._validate_gate(gate)
-        if control_qubit is not None:
-            operator = Gates.create_controlled_gate(
-                gate, control_qubit, target_qubit, self.num_qubits
-            )
-        else:
-            operator = np.eye(1)
-            for qubit in range(self.num_qubits):
-                operator = np.kron(
-                    operator,
-                    gate if qubit == target_qubit else np.eye(2),
-                )
-        self.state_vector = operator @ self.state_vector
-        self.circuit.append((gate_name, target_qubit, control_qubit))
-
-    def h(self, target_qubit: int):
-        """
-        Applies Hadamard gate to the target qubit.
-
-        :param target_qubit: Index of the target qubit.
-        """
-        self._apply_gate("H", Gates.H, target_qubit)
-
-    def t(self, target_qubit: int):
-        """
-        Applies π/8 gate to the target qubit.
-
-        :param target_qubit: Index of the target qubit.
-        """
-        self._apply_gate("T", Gates.T, target_qubit)
-
-    def x(self, target_qubit: int):
-        """
-        Applies Not gate to the target qubit.
-
-        :param target_qubit: Index of the target qubit.
-        """
-        self._apply_gate("X", Gates.X, target_qubit)
-
-    def cx(self, control_qubit: int, target_qubit: int):
-        """
-        Applies Controlled-Not gate to the target qubit.
-
-        :param control_qubit: Index of the control qubit.
-        :param target_qubit: Index of the target qubit.
-        """
-        self._apply_gate("X", Gates.X, target_qubit, control_qubit)
-
-    def u(
-        self,
-        target_qubit: int,
-        theta: float,
-        phi: float,
-        lambda_: float,
-        inverse: Optional[bool] = False,
-    ):
-        """
-        Applies Generic gate to the target qubit.
-
-        :param target_qubit: Index of the target qubit.
-        :param theta: Angle theta.
-        :param phi: Angle phi.
-        :param lambda_: Angle lambda.
-        :param inverse: Whether to apply the inverse of the gate.
-        """
-        gate = (
-            Gates.create_inverse_gate(Gates.U(theta, phi, lambda_))
-            if inverse
-            else Gates.U(theta, phi, lambda_)
-        )
+        self.n = num_qubits
+        # Statevector of length 2^n, start in |0...0>
+        self.state = np.zeros(2**num_qubits, dtype=complex)
+        self.state[0] = 1.0
+
+    def _apply_gate(self, U: np.ndarray, qubits: list):
+        k = len(qubits)  # number of qubits this gate acts on
+        shapeU = (2,) * k + (2,) * k  # e.g. for 2-qubit gate => (2,2, 2,2)
+        U_reshaped = U.reshape(shapeU)  # from (2^k,2^k) to (2,...,2,2,...,2)
+        st = self.state.reshape([2] * self.n)
+        # Move the targeted qubits' axes to the front, so we contract over them
+        st = np.moveaxis(st, qubits, range(k))
+        # tensordot over the last k dims of U with the first k dims of st
+        #   - The last k axes of U_reshaped are the "input" axes
+        #   - The first k axes of st are the qubits we apply the gate to
+        st_out = np.tensordot(U_reshaped, st, axes=(range(k, 2 * k), range(k)))
+        # st_out now has k "output" axes in front, plus the other (n-k) axes
+        # Move the front k axes back to their original positions
+        st_out = np.moveaxis(st_out, range(k), qubits)
+        # Flatten back to 1D
+        self.state = st_out.ravel()
+
+    # Single-qubit gates
+    def x(self, q: int):
+        self._apply_gate(Gates.X, [q])
+
+    def y(self, q: int):
+        self._apply_gate(Gates.Y, [q])
+
+    def z(self, q: int):
+        self._apply_gate(Gates.Z, [q])
+
+    def h(self, q: int):
+        self._apply_gate(Gates.H, [q])
+
+    def s(self, q: int):
+        self._apply_gate(Gates.S, [q])
+
+    def t(self, q: int):
+        self._apply_gate(Gates.T, [q])
+
+    def u(self, theta: float, phi: float, lam: float, q: int):
+        self._apply_gate(Gates.U(theta, phi, lam), [q])
+
+    # Two-qubit gates
+    def cx(self, control: int, target: int):
+        self._apply_gate(Gates.controlled_gate(Gates.X), [control, target])
+
+    def cu(self, theta: float, phi: float, lam: float, control: int, target: int):
         self._apply_gate(
-            self._get_gate_name(theta, phi, lambda_, inverse), gate, target_qubit
-        )
-
-    def cu(
-        self,
-        control_qubit: int,
-        target_qubit: int,
-        theta: float,
-        phi: float,
-        lambda_: float,
-        inverse: Optional[bool] = False,
-    ):
-        """
-        Applies Controlled-Generic gate to the target qubit.
-
-        :param control_qubit: Index of the control qubit.
-        :param target_qubit: Index of the target qubit.
-        :param theta: Angle theta.
-        :param phi: Angle phi.
-        :param lambda_: Angle lambda.
-        :param inverse: Whether to apply the inverse of the gate.
-        """
-        gate = (
-            Gates.create_inverse_gate(Gates.U(theta, phi, lambda_))
-            if inverse
-            else Gates.U(theta, phi, lambda_)
+            Gates.controlled_gate(Gates.U(theta, phi, lam)), [control, target]
         )
-        self._apply_gate(
-            self._get_gate_name(theta, phi, lambda_, inverse),
-            gate,
-            target_qubit,
-            control_qubit,
-        )
-
-    def measure(self, shots: int = 1, basis: Optional[np.ndarray] = None) -> List[str]:
-        """
-        Measures the state of the qubits.
-
-        :param shots: Number of measurements.
-        :param basis: Optional basis transformation.
-        :return: List of measurement results.
-        :raises ValueError: If the number of shots is negative.
-        """
-        if shots < 0:
-            raise ValueError("Number of shots must be non-negative.")
-        if basis is not None:
-            Gates._validate_gate(basis)
-            state_vector = basis @ self.state_vector
-        else:
-            state_vector = self.state_vector
-        probabilities = np.abs(state_vector) ** 2
-        counts = np.round(probabilities * shots).astype(int)
-        unique_states = [format(i, f"0{self.num_qubits}b") for i in range(len(counts))]
-        results = [
-            state for state, count in zip(unique_states, counts) for _ in range(count)
-        ]
-        diff = sum(counts) - shots
-        if diff > 0:
-            idx_to_remove = np.random.choice(len(results))
-            results.pop(idx_to_remove)
-        elif diff < 0:
-            idx_to_add = np.random.choice(len(counts), p=probabilities)
-            results.append(format(idx_to_add, f"0{self.num_qubits}b"))
-        return results
-
-    def run(
-        self, shots: int = 100, basis: Optional[np.ndarray] = None
-    ) -> Dict[str, int]:
-        """
-        Runs the simulation and returns measurement results.
-
-        :param shots: Number of measurements.
-        :param basis: Optional basis transformation.
-        :return: Dictionary of measurement results.
-        """
-        results = self.measure(shots, basis)
-        return dict(Counter(results))
-
-    def reset(self):
-        """
-        Resets the simulator to its initial state.
-        """
-        self.__init__(self.num_qubits)
-
-    def plot_wavefunction(self):
-        """
-        Plots the wavefunction's amplitude and phase using a phase circle plot.
-        """
-        amplitude = np.abs(self.state_vector)
-        phase = np.angle(self.state_vector)
-        labels = [
-            format(i, f"0{self.num_qubits}b") for i in range(len(self.state_vector))
-        ]
-        fig, ax = plt.subplots()
-        ax.set_aspect("equal", "box")
-        for i, (amp, phi) in enumerate(zip(amplitude, phase)):
-            x = amp * np.cos(phi)
-            y = amp * np.sin(phi)
-            ax.scatter(x, y)
-            ax.annotate(
-                labels[i],
-                (x, y),
-                textcoords="offset points",
-                xytext=(0, 10),
-                ha="center",
-            )
-        ax.set_xlim(-1.1, 1.1)
-        ax.set_ylim(-1.1, 1.1)
-        ax.axhline(0, color="black", linewidth=0.5)
-        ax.axvline(0, color="black", linewidth=0.5)
-        plt.title("Amplitude and Phase of Quantum States")
-        plt.xlabel("Real Component (Cosine of Phase * Amplitude)")
-        plt.ylabel("Imaginary Component (Sine of Phase * Amplitude)")
-        plt.show()
-
-    def __str__(self) -> str:
-        """
-        Returns a string representation of the circuit.
-
-        :return: String representing the circuit.
-        """
-        separator_length = sum(
-            (len(gate_name) + 3 for gate_name, _, _ in self.circuit), 1
-        )
-        lines = ["-" * separator_length]
-        qubit_lines = [["|"] for _ in range(self.num_qubits)]
-        for gate_name, target_qubit, control_qubit in self.circuit:
-            gate_name_length = len(gate_name)
-            gate_name_str = f" {gate_name} ".center(gate_name_length + 2, " ")
-            for i in range(self.num_qubits):
-                if control_qubit == i:
-                    qubit_lines[i].append(" @ ".center(gate_name_length + 2, " "))
-                elif target_qubit == i:
-                    qubit_lines[i].append(gate_name_str)
-                else:
-                    qubit_lines[i].append(" " * (gate_name_length + 2))
-                qubit_lines[i].append("|")
-        lines += ["".join(line) for line in qubit_lines]
-        lines += ["-" * separator_length]
-        return "\n".join(lines)
-
-    def __getsize__(self) -> int:
-        """
-        Returns the total memory size of the instance.
 
-        :return: Total memory size in bytes.
-        """
-        return getsizeof(self) + sum(map(getsizeof, self.__dict__.values()))
+    def swap(self, q1: int, q2: int):
+        self._apply_gate(Gates.SWAP_matrix(), [q1, q2])
+
+    def iswap(self, q1: int, q2: int):
+        self._apply_gate(Gates.iSWAP_matrix(), [q1, q2])
+
+    # Three-qubit gates
+    def toffoli(self, c1: int, c2: int, t: int):
+        self._apply_gate(Gates.Toffoli_matrix(), [c1, c2, t])
+
+    def fredkin(self, c: int, t1: int, t2: int):
+        self._apply_gate(Gates.Fredkin_matrix(), [c, t1, t2])
+
+    # Simulation & measurement
+    def run(self, shots: int = 100) -> dict[str, int]:
+        # Compute base counts by multiplying probabilities and truncating
+        float_counts = shots * (np.abs(self.state)**2)
+        base_counts = float_counts.astype(int)
+        remainder = shots - base_counts.sum()
+        if remainder:
+            # Distribute leftover shots to states with the largest fractional parts
+            fractional_parts = float_counts - base_counts
+            base_counts[np.argsort(fractional_parts)[-remainder:]] += 1
+        # Return only those states that actually occurred
+        return {f"{i:0{self.n}b}": int(c) for i, c in enumerate(base_counts) if c}
diff --git a/setup.py b/setup.py
index 2bd3a1e..fa69e50 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="qubit-simulator",
-    version="0.0.8",
+    version="0.0.9",
     description="A simple qubit simulator",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",
diff --git a/tests/test_gates.py b/tests/test_gates.py
index 4000b77..cde6cdb 100644
--- a/tests/test_gates.py
+++ b/tests/test_gates.py
@@ -1,32 +1,94 @@
-import pytest
 import numpy as np
 from qubit_simulator import Gates
 
 
-def test_create_inverse_gate():
-    random_matrix = np.random.rand(2, 2) + 1j * np.random.rand(2, 2)
-    random_unitary_gate, _ = np.linalg.qr(random_matrix)
-    inverse_gate = Gates.create_inverse_gate(random_unitary_gate)
-    assert np.allclose(random_unitary_gate @ inverse_gate, np.eye(2))
+def test_gate_shapes():
+    """Check that the gates have the correct matrix dimensions."""
+    # Single-qubit gates => 2x2
+    for gate in [Gates.X, Gates.Y, Gates.Z, Gates.H, Gates.S, Gates.T]:
+        assert gate.shape == (2, 2), f"Gate {gate} should be 2x2."
 
+    # SWAP, iSWAP => 4x4
+    assert Gates.SWAP_matrix().shape == (4, 4), "SWAP should be 4x4."
+    assert Gates.iSWAP_matrix().shape == (4, 4), "iSWAP should be 4x4."
 
-def test_create_controlled_gate():
-    # Test for a controlled-X gate
-    controlled_X = Gates.create_controlled_gate(Gates.X, 0, 1, 2)
-    expected_result = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]])
-    assert np.allclose(controlled_X, expected_result)
+    # Toffoli, Fredkin => 8x8
+    assert Gates.Toffoli_matrix().shape == (8, 8), "Toffoli should be 8x8."
+    assert Gates.Fredkin_matrix().shape == (8, 8), "Fredkin should be 8x8."
 
 
-def test_non_unitary_gate():
-    # Define a non-unitary matrix (conjugate transpose is not equal to its inverse)
-    non_unitary_gate = np.array([[2, 0], [0, 0.5]])
-    with pytest.raises(ValueError):
-        Gates._validate_gate(non_unitary_gate)
+def test_is_unitary_single_qubit():
+    """Check that single-qubit gates are unitary (U * U^† = I)."""
+    for gate in [Gates.X, Gates.Y, Gates.Z, Gates.H, Gates.S, Gates.T]:
+        U_dagger = gate.conjugate().T
+        identity = np.dot(gate, U_dagger)
+        assert np.allclose(identity, np.eye(2, dtype=complex)), f"{gate} not unitary."
 
 
-def test_create_controlled_gate_invalid_qubits():
-    # Define a scenario where the control and target qubits are out of range
-    with pytest.raises(IndexError):
-        Gates.create_controlled_gate(
-            Gates.X, control_qubit=4, target_qubit=5, num_qubits=3
-        )
+def test_is_unitary_two_qubit():
+    """Check that two-qubit gates are unitary."""
+    two_qubit_gates = [Gates.SWAP_matrix(), Gates.iSWAP_matrix()]
+    for gate in two_qubit_gates:
+        U_dagger = gate.conjugate().T
+        identity = np.dot(gate, U_dagger)
+        assert np.allclose(identity, np.eye(4, dtype=complex)), f"{gate} not unitary."
+
+
+def test_is_unitary_three_qubit():
+    """Check that three-qubit gates are unitary."""
+    three_qubit_gates = [Gates.Toffoli_matrix(), Gates.Fredkin_matrix()]
+    for gate in three_qubit_gates:
+        U_dagger = gate.conjugate().T
+        identity = np.dot(gate, U_dagger)
+        assert np.allclose(identity, np.eye(8, dtype=complex)), f"{gate} not unitary."
+
+
+def test_controlled_gate_shape():
+    """Check that controlled_gate(U) is 4x4 if U is 2x2."""
+    for gate in [Gates.X, Gates.Y, Gates.Z, Gates.H, Gates.S, Gates.T]:
+        c_gate = Gates.controlled_gate(gate)
+        assert c_gate.shape == (4, 4), "Controlled version of a 2x2 gate should be 4x4."
+
+
+def test_controlled_gate_blocks():
+    """Check that controlled_gate structure matches block-diagonal form: I (2x2) on top-left, U on bottom-right."""
+    U = Gates.X  # any 2x2
+    c_gate = Gates.controlled_gate(U)
+    # top-left 2x2 should be identity
+    assert np.allclose(c_gate[:2, :2], np.eye(2)), "Top-left block should be identity."
+    # bottom-right 2x2 should be U
+    assert np.allclose(
+        c_gate[2:, 2:], U
+    ), "Bottom-right block should be the original gate."
+
+
+def test_inverse_gate_single_qubit():
+    """Check that inverse_gate(U) = U^† for single-qubit gates."""
+    for gate in [Gates.X, Gates.Y, Gates.Z, Gates.H, Gates.S, Gates.T]:
+        U_inv = Gates.inverse_gate(gate)
+        U_dagger = gate.conjugate().T
+        assert np.allclose(U_inv, U_dagger), "inverse_gate should match U^†."
+
+
+def test_U_parameterized_gate():
+    """
+    Check dimension & a known special case for U(θ, φ, λ).
+    We'll use angles that produce X up to a global phase.
+    """
+    import math
+
+    # 1) Basic shape check
+    Umat = Gates.U(theta=math.pi / 2, phi=0, lam=0)
+    assert Umat.shape == (2, 2), "Parameterized U gate should be 2x2."
+
+    # 2) Produce X up to global phase with U(π, -π/2, π/2)
+    U_x = Gates.U(theta=math.pi, phi=-math.pi / 2, lam=math.pi / 2)
+    X = Gates.X
+
+    # Compare up to a global phase
+    # We'll pick a nonzero element to find the ratio
+    ratio = U_x[0, 1] / X[0, 1]  # e.g., compare top-right
+    adjusted = U_x * np.conjugate(ratio)
+    assert np.allclose(
+        adjusted, X, atol=1e-8
+    ), "U(π, -π/2, π/2) not matching X up to a global phase."
diff --git a/tests/test_simulator.py b/tests/test_simulator.py
index 43e48f0..1d7b91e 100644
--- a/tests/test_simulator.py
+++ b/tests/test_simulator.py
@@ -1,391 +1,225 @@
 import pytest
 import numpy as np
-import matplotlib.pyplot as plt
 from qubit_simulator import QubitSimulator, Gates
 
-# Initialization and Basic Configuration
 
+# 1. Initialization Tests
+def test_initial_state_vector_length():
+    sim = QubitSimulator(num_qubits=3)
+    assert len(sim.state) == 2**3, "State vector length should be 2^n."
 
-def test_initial_state():
-    simulator = QubitSimulator(3)
-    assert np.allclose(simulator.state_vector, [1, 0, 0, 0, 0, 0, 0, 0])
 
+def test_initial_state_is_zero_state():
+    sim = QubitSimulator(num_qubits=3)
+    expected = np.zeros(2**3, dtype=complex)
+    expected[0] = 1.0
+    assert np.allclose(sim.state, expected), "Initial state should be |000>."
 
-def test_large_number_of_qubits():
-    num_qubits = 20
-    simulator = QubitSimulator(num_qubits)
-    assert len(simulator.state_vector) == 2**num_qubits
 
+# 2. Single-Qubit Gate Tests
+def test_x_gate_on_single_qubit():
+    sim = QubitSimulator(num_qubits=1)
+    sim.x(0)
+    expected = np.array([0, 1], dtype=complex)
+    assert np.allclose(sim.state, expected), "X gate did not produce |1> from |0>."
 
-def test_zero_qubits():
-    simulator = QubitSimulator(0)
-    assert len(simulator.state_vector) == 1
-    assert simulator.state_vector[0] == 1
 
+def test_y_gate_on_single_qubit():
+    sim = QubitSimulator(num_qubits=1)
+    sim.y(0)
+    expected = np.array([0, 1j], dtype=complex)  # i|1>
+    assert np.allclose(sim.state, expected), "Y gate did not produce i|1> from |0>."
 
-def test_negative_qubits():
-    with pytest.raises(ValueError):
-        QubitSimulator(-1)  # Negative number of qubits is not allowed
 
+def test_z_gate_on_single_qubit():
+    sim = QubitSimulator(num_qubits=1)
+    sim.x(0)  # now |1>
+    sim.z(0)  # => -|1>
+    expected = np.array([0, -1], dtype=complex)
+    assert np.allclose(sim.state, expected), "Z gate did not produce -|1>."
 
-def test_initialization_complex_states():
-    simulator = QubitSimulator(2)
-    simulator.state_vector = [0.5, 0.5, 0.5, 0.5]
-    simulator.x(0)
-    assert np.allclose(simulator.state_vector, [0.5, 0.5, 0.5, 0.5])
 
-
-# Gate Operations (Single Qubit)
+def test_h_gate_on_single_qubit():
+    sim = QubitSimulator(num_qubits=1)
+    sim.h(0)
+    expected = np.array([1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=complex)
+    assert np.allclose(
+        sim.state, expected
+    ), "H gate should produce (|0> + |1>)/sqrt(2)."
 
 
-def test_x_gate():
-    simulator = QubitSimulator(1)
-    simulator.x(0)
-    # After applying the Pauli-X gate, the state should be |1⟩
-    assert np.allclose(simulator.state_vector, [0, 1])
+def test_h_gate_twice():
+    sim = QubitSimulator(num_qubits=1)
+    sim.h(0)
+    sim.h(0)
+    expected = np.array([1, 0], dtype=complex)
+    assert np.allclose(sim.state, expected), "H applied twice should return to |0>."
 
 
-def test_h_gate():
-    simulator = QubitSimulator(1)
-    simulator.h(0)
-    # After applying the Hadamard gate, the state should be an equal superposition
-    assert np.allclose(simulator.state_vector, [0.70710678, 0.70710678])
+def test_s_gate():
+    sim = QubitSimulator(num_qubits=1)
+    sim.x(0)  # |1>
+    sim.s(0)  # => i|1>
+    expected = np.array([0, 1j], dtype=complex)
+    assert np.allclose(sim.state, expected), "S gate did not apply phase i to |1>."
 
 
 def test_t_gate():
-    simulator = QubitSimulator(1)
-    simulator.x(0)  # Set the initial state to |1⟩
-    simulator.t(0)
-    # After applying the π/8 gate, the state should have a phase shift of π/4
-    assert np.allclose(simulator.state_vector, [0, 0.70710678 + 0.70710678j])
+    sim = QubitSimulator(num_qubits=1)
+    sim.x(0)  # |1>
+    sim.t(0)  # => e^{i pi/4}|1>
+    phase = np.exp(1j * np.pi / 4)
+    expected = np.array([0, phase], dtype=complex)
+    assert np.allclose(sim.state, expected), "T gate did not produce e^{i pi/4}|1>."
 
 
 def test_u_gate():
-    theta = np.pi / 4
-    phi = np.pi / 3
-    lambda_ = np.pi / 2
-    simulator = QubitSimulator(1)
-    simulator.u(0, theta, phi, lambda_)
-    # Expected result obtained from the U matrix using the given parameters
-    expected_result = Gates.U(theta, phi, lambda_) @ [1, 0]
-    assert np.allclose(simulator.state_vector, expected_result)
-
-
-@pytest.mark.parametrize(
-    "theta,phi,lambda_", [(0, 0, 0), (2 * np.pi, 2 * np.pi, 2 * np.pi)]
-)
-def test_u_gate_edge_cases(theta, phi, lambda_):
-    simulator = QubitSimulator(1)
-    simulator.u(0, theta, phi, lambda_)
-    # State vector should be |0⟩
-    assert np.allclose(simulator.state_vector, [1, 0])
-
-
-# Gate Operations (Multi-Qubit)
-
-
-def test_cx_gate():
-    simulator = QubitSimulator(2)
-    simulator.state_vector = [0, 0, 0, 1]  # Set the initial state to |11⟩
-    simulator.cx(0, 1)
-    # After applying the CNOT gate, the state should be |10⟩ (big-endian)
-    assert np.allclose(simulator.state_vector, [0, 0, 1, 0])
-
-
-def test_cu_gate():
-    theta = np.pi / 4
-    phi = np.pi / 3
-    lambda_ = np.pi / 2
-    simulator = QubitSimulator(2)
-    simulator.x(0)  # Set the control qubit to |1⟩
-    simulator.cu(0, 1, theta, phi, lambda_)
-    # Initial state |10⟩
-    initial_state = np.array([0, 0, 1, 0], dtype=complex)
-    # Apply U gate to the target qubit
-    expected_result = np.kron(np.eye(2), Gates.U(theta, phi, lambda_)) @ initial_state
-    assert np.allclose(simulator.state_vector, expected_result)
-
-
-def test_cu_gate_no_effect():
-    theta = np.pi / 4
-    phi = np.pi / 3
-    lambda_ = np.pi / 2
-    simulator = QubitSimulator(2)
-    # Control qubit is |0⟩, so the CU gate should have no effect
-    simulator.cu(0, 1, theta, phi, lambda_)
-    assert np.allclose(simulator.state_vector, [1, 0, 0, 0])
-
-
-def test_target_control():
-    simulator = QubitSimulator(3)
-    simulator.x(0)
-    simulator.x(2)  # Set the initial state to |101⟩
-    simulator.cx(control_qubit=2, target_qubit=0)
-    # After applying the CNOT gate, the state should be |001⟩
-    assert np.allclose(simulator.state_vector, [0, 1, 0, 0, 0, 0, 0, 0])
-
-
-# Measurement and Probabilities
-
-
-def test_measure():
-    simulator = QubitSimulator(1)
-    simulator.x(0)
-    # After applying the X gate, the state should be |1⟩
-    result = simulator.measure()
-    assert result == ["1"]
-
-
-@pytest.mark.parametrize("shots", [2, 10])
-def test_measure_multiple_shots(shots):
-    simulator = QubitSimulator(1)
-    simulator.x(0)
-    results = simulator.measure(shots=shots)
-    assert results.count("1") == shots
-
-
-@pytest.mark.parametrize("shots", [5, 13])
-def test_measure_adjustment(shots):
-    simulator = QubitSimulator(3)
-    simulator.h(0)
-    simulator.t(1)
-    simulator.cx(0, 2)
-    simulator.u(2, 0.3, 0.4, 0.5)
-    results = simulator.measure(shots)
-    assert len(results) == shots
-
-
-@pytest.mark.parametrize("shots", [-1, -10])
-def test_negative_shots(shots):
-    simulator = QubitSimulator(1)
-    with pytest.raises(ValueError):
-        simulator.run(shots=shots)  # Negative shots are invalid
-
-
-def test_measure_without_gates():
-    simulator = QubitSimulator(2)
-    results = simulator.run(shots=100)
-    assert results == {"00": 100}
-
-
-def test_measure_custom_basis():
-    simulator = QubitSimulator(1)
-    # Define the transformation matrix for the Pauli-X basis
-    X_basis = np.array([[0, 1], [1, 0]])
-    # Apply the X gate to the qubit, transforming it to |1⟩
-    simulator.x(0)
-    # Measure in the X basis, which should result in the state |0⟩ in the X basis
-    result = simulator.run(shots=10, basis=X_basis)
-    assert set(result) == {"0"}
-
+    sim = QubitSimulator(num_qubits=1)
+    sim.u(theta=np.pi, phi=0, lam=0, q=0)  # acts like X on |0> (up to global phase)
+    expected = np.array([0, 1], dtype=complex)
+    assert np.allclose(sim.state, expected), "U(π,0,0) not matching X action on |0>."
+
+
+# 3. Two-Qubit Gate Tests
+def test_cnot_control_zero():
+    sim = QubitSimulator(num_qubits=2)
+    # |00>, apply CNOT => still |00>
+    sim.cx(0, 1)
+    expected = np.array([1, 0, 0, 0], dtype=complex)
+    assert np.allclose(sim.state, expected), "CNOT changed target when control was |0>."
+
+
+def test_cnot_control_one():
+    sim = QubitSimulator(num_qubits=2)
+    sim.x(0)  # now |10>
+    sim.cx(0, 1)  # => |11>
+    expected = np.array([0, 0, 0, 1], dtype=complex)
+    assert np.allclose(
+        sim.state, expected
+    ), "CNOT did not flip target when control was |1>."
 
-def test_measure_custom_basis_valid():
-    simulator = QubitSimulator(1)
-    Z_basis = np.array([[1, 0], [0, -1]])
-    simulator.x(0)
-    result = simulator.measure(basis=Z_basis)
-    assert result == ["1"]
 
+def test_swap_gate():
+    sim = QubitSimulator(num_qubits=2)
+    sim.x(1)  # => |01>
+    sim.swap(0, 1)  # => |10>
+    expected = np.array([0, 0, 1, 0], dtype=complex)
+    assert np.allclose(sim.state, expected), "SWAP did not swap |01> to |10>."
 
-def test_measure_probabilities():
-    shots = 1000
-    simulator = QubitSimulator(1)
-    simulator.h(0)
-    results = simulator.run(shots=shots)
-    assert results.get("0", 0) == results.get("1", 0)
-
-
-# Error Handling and Validation
-
-
-def test_invalid_basis_transformation():
-    simulator = QubitSimulator(1)
-    # Define an invalid basis transformation (not unitary)
-    invalid_basis = np.array([[1, 2], [2, 1]])
-    with pytest.raises(ValueError):
-        simulator.run(basis=invalid_basis)
 
+def test_iswap_gate():
+    sim = QubitSimulator(num_qubits=2)
+    sim.x(1)  # => |01>
+    sim.iswap(0, 1)  # => i|10>
+    expected = np.array([0, 0, 1j, 0], dtype=complex)
+    assert np.allclose(sim.state, expected), "iSWAP did not produce i|10> from |01>."
 
-def test_invalid_qubit_index():
-    simulator = QubitSimulator(1)
-    with pytest.raises(IndexError):
-        simulator.h(2)  # Index out of range
-
-
-def test_reset_invalid_qubit_index():
-    simulator = QubitSimulator(3)
-    simulator.num_qubits = -1  # Set an invalid value for num_qubits
-    with pytest.raises(ValueError):
-        simulator.reset()  # Resetting with an invalid value should raise an error
-
-
-def test_invalid_control_and_target_index():
-    simulator = QubitSimulator(1)
-    with pytest.raises(IndexError):
-        simulator.cx(1, 0)  # Control qubit cannot be out of range
-
-
-def test_apply_gate_invalid_control_qubit():
-    simulator = QubitSimulator(1)
-    with pytest.raises(IndexError):
-        simulator._apply_gate("X", Gates.X, target_qubit=0, control_qubit=2)
-
-
-def test_error_messages():
-    with pytest.raises(ValueError, match="Number of qubits must be non-negative."):
-        QubitSimulator(-1)
-    with pytest.raises(ValueError, match="Number of shots must be non-negative."):
-        QubitSimulator(1).measure(-1)
-
-
-# Circuit Functionality
-
-
-def test_circuit_reset():
-    simulator = QubitSimulator(1)
-    simulator.x(0)
-    simulator.reset()
-    assert np.allclose(simulator.state_vector, [1, 0])
-
-
-def test_run():
-    simulator = QubitSimulator(1)
-    # Running the simulation 10 times should produce 10 results
-    results = simulator.run(10)
-    assert results == {"0": 10}
 
-
-def test_gate_reversibility():
-    theta = np.pi / 2
-    phi = np.pi / 4
-    lambda_ = np.pi / 3
-    simulator = QubitSimulator(1)
-    simulator.u(0, theta, phi, lambda_)
-    simulator.h(0)
-    simulator.x(0)
-    simulator.x(0)
-    simulator.h(0)
-    # Apply U inverse
-    simulator.u(0, theta, phi, lambda_, inverse=True)
-    assert np.allclose(simulator.state_vector, [1, 0])
-
-
-def test_random_unitary_gate_inverse():
-    simulator = QubitSimulator(1)
-    # Generate a random unitary matrix using QR decomposition
-    random_matrix = np.random.rand(2, 2) + 1j * np.random.rand(2, 2)
-    random_unitary_gate, _ = np.linalg.qr(random_matrix)
-    simulator._apply_gate("RANDOM_UNITARY", random_unitary_gate, 0)
-    simulator._apply_gate("RANDOM_UNITARY_INV", random_unitary_gate.conj().T, 0)
-    # The final state should be the same as the initial state
-    assert np.allclose(simulator.state_vector, [1, 0])
-
-
-@pytest.mark.parametrize(
-    "num_qubits, expected_string",
-    [
-        (0, "-\n-"),
-        (
-            3,
-            "-----------------------------------\n"
-            "| H |   |   |          @          |\n"
-            "|   | X | X | U(1.05, 0.63, 0.45) |\n"
-            "|   |   | @ |                     |\n"
-            "-----------------------------------",
-        ),
-    ],
-)
-def test_circuit_string(num_qubits, expected_string):
-    simulator = QubitSimulator(num_qubits)
-    if num_qubits == 3:
-        simulator.h(0)
-        simulator.x(1)
-        simulator.cx(2, 1)
-        simulator.cu(0, 1, np.pi / 3, np.pi / 5, np.pi / 7)
-    assert str(simulator) == expected_string
-
-
-def test_complex_circuit():
-    simulator = QubitSimulator(3)
-    simulator.h(0)
-    simulator.u(1, np.pi / 4, np.pi / 4, np.pi / 2)
-    simulator.cx(2, 0)
-    simulator.cu(1, 2, np.pi / 2, np.pi / 4, np.pi / 8)
-    simulator.x(0)
-    simulator.run(shots=10)
-    # This test verifies the process rather than the final state, so no assertion is needed
+def test_controlled_u_gate():
+    sim = QubitSimulator(num_qubits=2)
+    # control=0 => not activated => remains |00>
+    sim.cu(theta=np.pi, phi=0, lam=0, control=0, target=1)
+    expected = np.array([1, 0, 0, 0], dtype=complex)
+    assert np.allclose(
+        sim.state, expected
+    ), "Controlled-U acted even though control was |0>."
+
+    # Now put control in |1> => |10>, apply CU => acts like X on target => |11>
+    sim = QubitSimulator(num_qubits=2)
+    sim.x(0)  # => |10>
+    sim.cu(theta=np.pi, phi=0, lam=0, control=0, target=1)
+    expected = np.array([0, 0, 0, 1], dtype=complex)
+    assert np.allclose(sim.state, expected), "Controlled-U didn't act like X on |10>."
+
+
+# 4. Three-Qubit Gate Tests
+def test_toffoli_gate():
+    sim = QubitSimulator(num_qubits=3)
+    sim.x(0)
+    sim.x(1)  # => |110>
+    sim.toffoli(0, 1, 2)  # => should flip target => |111>
+    expected = np.zeros(8, dtype=complex)
+    expected[7] = 1
+    assert np.allclose(sim.state, expected), "Toffoli did not flip target for |110>."
+
+
+def test_toffoli_partial_control():
+    sim = QubitSimulator(num_qubits=3)
+    sim.x(0)  # => |100>, only 1st ctrl is 1 => no flip => remains |100>
+    sim.toffoli(0, 1, 2)
+    expected = np.zeros(8, dtype=complex)
+    expected[4] = 1
+    assert np.allclose(
+        sim.state, expected
+    ), "Toffoli flipped target even though second ctrl was 0."
+
+
+def test_fredkin_gate():
+    # Fredkin = CSWAP: if ctrl=1 => swap qubits 1 & 2
+    sim = QubitSimulator(num_qubits=3)
+    sim.x(0)  # qubit0=1
+    sim.x(2)  # => |101>
+    sim.fredkin(0, 1, 2)  # => swap qubits 1 & 2 => |110>
+    expected = np.zeros(8, dtype=complex)
+    expected[6] = 1
+    assert np.allclose(
+        sim.state, expected
+    ), "Fredkin did not swap qubits 1 & 2 correctly."
+
+
+# 5. Simulator-Specific Tests (Measurement, Norm, etc.)
+def test_apply_gate_then_inverse():
+    """
+    Use the simulator to apply a gate then its inverse,
+    verifying the final state returns to |0>.
+    """
+    sim = QubitSimulator(num_qubits=1)
+    sim.h(0)
+    sim._apply_gate(Gates.inverse_gate(Gates.H), [0])
+    expected = np.array([1, 0], dtype=complex)
+    assert np.allclose(
+        sim.state, expected
+    ), "Applying H then H^† did not return to |0>."
+
+
+def test_measurement_counts_no_superposition():
+    sim = QubitSimulator(num_qubits=2)
+    # => |00> by default
+    shots = 200
+    counts = sim.run(shots=shots)
+    assert (
+        counts.get("00", 0) == shots
+    ), "All measurements should be '00' in a definite state."
+
+
+@pytest.mark.parametrize("shots", [1000, 2000])
+def test_measurement_distribution(shots):
+    sim = QubitSimulator(num_qubits=1)
+    sim.h(0)
+    counts = sim.run(shots=shots)
+    zero_counts = counts.get("0", 0)
+    one_counts = counts.get("1", 0)
+    # Expect roughly half 0, half 1
+    assert abs(zero_counts - shots / 2) < 5 * np.sqrt(
+        shots
+    ), "Distribution off from 50% for '0'."
+    assert abs(one_counts - shots / 2) < 5 * np.sqrt(
+        shots
+    ), "Distribution off from 50% for '1'."
 
 
 def test_bell_state():
-    simulator = QubitSimulator(2)
-    simulator.h(0)
-    simulator.cx(0, 1)
-    # After applying the Hadamard and CNOT gates, the state should be a Bell state
-    assert np.allclose(simulator.state_vector, [0.70710678, 0, 0, 0.70710678])
-
-
-def test_ghz_state():
-    simulator = QubitSimulator(3)
-    simulator.h(0)
-    simulator.cx(0, 1)
-    simulator.cx(0, 2)
-    # After applying the Hadamard and CNOT gates, the state should be a GHZ state
-    assert np.allclose(
-        simulator.state_vector, [0.70710678, 0, 0, 0, 0, 0, 0, 0.70710678]
-    )
-
-
-@pytest.mark.parametrize("num_qubits", [1, 2, 5])
-def test_qft(num_qubits):
-    def apply_qft(simulator):
-        num_qubits = simulator.num_qubits
-        for target_qubit in range(num_qubits):
-            simulator.h(target_qubit)
-            for control_qubit in range(target_qubit + 1, num_qubits):
-                phase_angle = 2 * np.pi / (2 ** (control_qubit - target_qubit + 1))
-                simulator.cu(control_qubit, target_qubit, 0, -phase_angle, 0)
-        # Swap qubits to match the desired output order
-        for i in range(num_qubits // 2):
-            j = num_qubits - i - 1
-            simulator.cx(i, j)
-            simulator.cx(j, i)
-            simulator.cx(i, j)
-
-    simulator = QubitSimulator(num_qubits)
-    # Create a random initial state vector and normalize it
-    random_state = np.random.rand(2**num_qubits) + 1j * np.random.rand(
-        2**num_qubits
-    )
-    random_state /= np.linalg.norm(random_state)
-    # Set the random state as the initial state in the simulator
-    simulator.state_vector = random_state.copy()
-    # Apply QFT in the simulator
-    apply_qft(simulator)
-    # Compute the expected result using NumPy's FFT and normalize
-    fft_result = np.fft.fft(random_state) / np.sqrt(2**num_qubits)
-    # Compare the state vectors
-    assert np.allclose(simulator.state_vector, fft_result)
-
-
-# Memory and Object Size
-
-
-def test_getsize():
-    simulator = QubitSimulator(2)
-    initial_size = simulator.__getsize__()
-    simulator.h(0)
-    simulator.cx(0, 1)
-    assert simulator.__getsize__() > initial_size
-
-
-# Plotting
-
-
-def test_plot_wavefunction():
-    simulator = QubitSimulator(3)
-    simulator.h(0)
-    simulator.t(1)
-    simulator.cx(0, 2)
-    simulator.u(2, 0.3, 0.4, 0.5)
-    simulator.plot_wavefunction()
-    plt.close("all")
+    sim = QubitSimulator(num_qubits=2)
+    sim.h(0)
+    sim.cx(0, 1)
+    expected = np.array([1 / np.sqrt(2), 0, 0, 1 / np.sqrt(2)], dtype=complex)
+    assert np.allclose(sim.state, expected), "Bell state not correct."
+
+
+def test_state_norm_is_one():
+    sim = QubitSimulator(num_qubits=2)
+    sim.h(0)
+    sim.s(0)
+    sim.cx(0, 1)
+    norm = np.sum(np.abs(sim.state) ** 2)
+    assert np.isclose(norm, 1.0), "State norm deviated from 1."