input_python.txt

================================================================================
  implementation/python/  --  PyTorch tensor-ops quantum simulator (Python 3.12+)
================================================================================


================================================================================
FILE: implementation/python/README.md
================================================================================
# implementation/python -- PyTorch state-vector quantum simulator

A pure-Python state-vector quantum-circuit simulator using PyTorch for
tensor ops and device-agnostic GPU acceleration. Sibling of
`implementation/c` (MPI) and `implementation/go` (goroutines), covering
the same thesis claims via a third parallelism model: tensor ops on
whatever device PyTorch can reach (NVIDIA CUDA, AMD ROCm, Apple Metal /
MPS), with a CPU fallback.

## Status

**Feature-complete.** The implementation covers the same thesis claims
as `/c` and `/go`:

- `qubit.Qreg` — state-vector register with construction, accessors,
  the qubit-axis helper, memory preflight, and a CPU measurement RNG.
- `qubit.standart` — `gcd_u64`, `mod_pow`, `continued_fraction`,
  `is_power_of_two`, `ilog2_u32`. Python builtins (`math.gcd`, `pow`,
  `fractions.Fraction.limit_denominator`) wrapped with
  `qubit:`-prefixed validation.
- Single-qubit gates — `apply_u` (the `tensordot` + `movedim`
  workhorse) plus the named gates `apply_h`, `apply_x`, `apply_y`,
  `apply_z`, `apply_s`, `apply_t`, `apply_phase`, `apply_rx`,
  `apply_ry`, `apply_rz`.
- Controlled gates — `apply_cu` (4x4 block-diagonal workhorse via
  permute + matmul + inverse-permute) plus `apply_cnot`, `apply_cz`,
  `apply_controlled_phase`, `apply_swap` (the 3-CNOT identity).
- Multi-controlled gates — `apply_multi_controlled_z` for phase-flip
  diffusion and `apply_multi_controlled_x` (generalised Toffoli),
  both vectorized mask/gather/scatter; no per-amplitude loops.
- Measurement — `measure_qubit` (single-qubit projective with collapse
  and renormalise), `measure_all` (sample full basis from `|amp|^2`),
  `sample_distribution` (snapshot + restore so the original isn't
  mutated), `clone` (independent amp + RNG state), `dump` (structured
  list of non-zero amplitudes).
- QFT — `apply_qft` (forward, includes final bit-reversal swaps so
  output is in natural binary order) and `apply_qft_inverse` (true
  inverse via reversed loop order and negated phase angles, not three
  forward applications). Sub-register support via `start=` / `n=`.
- Grover — `apply_grover(q, n_qubits, oracle, user=None,
  iterations=None)`. The oracle is a `Callable[[Qreg, Any], None]`
  that phase-flips marked basis states; the default iteration count
  is the optimum for one marked item (`floor(pi/4 * sqrt(N))`).
- Shor — `apply_modular_exp` (vectorized permutation via
  `torch.gather` with a CPU-built index tensor), `apply_shor_period`
  (period-finding circuit), `shor_factor(N, max_attempts=20,
  seed=None)` (end-to-end factoring, bit-for-bit reproducible when
  seeded). v1 covers Shor-15 (13-qubit register: `n=4` target +
  `t=2n+1=9` counting) and Shor-21 (16-qubit register: `n=5` + `t=11`,
  gated by `RUN_SHOR_21=1` to keep the default test loop fast).
- `qubit-demo` CLI — `--algo {bell,qft,grover,shor}` for quick
  algorithm-level smoke tests. Installed via `[project.scripts]`.

Function-style (`apply_h(q, 0)`) and method-style (`q.apply_h(0)`)
call shapes are both first-class for every gate, measurement op, QFT,
Grover, and Shor primitive.

**Tests:** the exact collected count is device-dependent because the
device-parametrised tests in `tests/conftest.py` expand once per
available backend. On the author's CPU + MPS host the count is
**414 default + 1 gated Shor-21 = 415 total**; on CPU-only CI runners
the MPS rows drop and the collected count is lower. All green on
every available device under `ruff check`, `mypy --strict`, and
pytest.

## Quickstart

The Makefile targets assume [`uv`](https://docs.astral.sh/uv/) is on
your path (it manages the venv and dev-tool installation). Without
`uv`, the same workflow runs through `pip install -e .[dev]` and your
preferred environment manager.

```bash
cd implementation/python
make sync                      # creates .venv, installs torch + dev tools
make test                      # ~0.5s on Apple Silicon
make check                     # lint + typecheck + test (the full PR gate)

make demo                      # default: bell
make demo ALGO=qft             # QFT|0> = uniform
make demo ALGO=grover          # Grover marks |1111>
make demo ALGO=shor            # factor 15

RUN_SHOR_21=1 uv run pytest tests/test_shor.py -q   # 16-qubit gated test
```

The CLI also runs directly: `uv run qubit-demo --algo shor`.

On many-core CPU-only hosts, PyTorch's default thread fan-out can be
counterproductive for the small tensor kernels these tests exercise.
If the gated Shor-21 path is slow or hangs, cap the thread count:

```bash
OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
    RUN_SHOR_21=1 uv run pytest tests/test_shor.py -q
```

Apple Silicon / MPS does not need this.

## Key design choices

- **Device-agnostic via `torch.device`.** `Qreg(n_qubits)` auto-detects
  `cuda` > `mps` > `cpu`. Override with `Qreg(n_qubits, device='cpu')`.
- **Tensor-native gates.** State is a flat `(2**n,)` complex tensor;
  gate code reshapes into the `(2,) * n` view and `einsum`/`tensordot`
  the unitary against the target axis. No per-amplitude loops.
- **Dtype policy.** `complex128` on CPU / CUDA / ROCm; `complex64` on
  MPS (the MPS backend lacks `float64`, so `complex128` is impossible
  there). Explicit `Qreg(..., device='mps', dtype=torch.complex128)`
  raises `ValueError` with a guidance message; the simulator never
  silently downgrades.
- **Qubit 0 is LSB.** Matches `/c` and `/go`. The tensor axis for
  qubit `q` is `n_qubits - 1 - q`, centralised in `qubit/_axis.py`.
- **CPU measurement RNG.** Even when `_amp` lives on GPU, the
  measurement RNG is a CPU `torch.Generator` — MPS-generator quirks
  are real and measurement is a readout boundary anyway. Seeded tests
  are bit-identical across CPU / CUDA / MPS.
- **No `destroy()` / `close()` / context manager.** Python's GC
  reclaims PyTorch tensors when the `Qreg` goes out of scope. No
  `with` block, no manual cleanup.

## API at a glance

```python
import math, torch
from qubit import Qreg, apply_h, apply_x, qubit_axis, shor_factor

# Auto-detect device + dtype.
q = Qreg(n_qubits=4)
print(q.device, q.dtype, q.n_qubits)

# Or explicit:
q = Qreg(4, device='cpu', seed=42, dtype=torch.complex128)

# Seeded measurement-RNG is reproducible across same-seed Qregs.
q.init_basis(5)  # |0101>
print(q.prob_of(5))      # 1.0
print(q.norm())          # 1.0

# Gates: both function-style and method-style work.
q.apply_h(0)             # method-style
apply_x(q, 1)            # function-style
print(q.norm())          # still 1.0 (unitary)

# Custom unitary via apply_u (the workhorse every named gate dispatches to).
inv2 = 1.0 / math.sqrt(2.0)
h = torch.tensor(
    [[inv2 + 0j, inv2 + 0j], [inv2 + 0j, -inv2 + 0j]],
    dtype=q.dtype, device=q.device,
)
q.apply_u(2, h)

# Defensive CPU clone for inspection / cross-implementation comparison:
amps = q.amplitudes_copy()   # 1-D CPU tensor, len 16

# Qubit-to-axis helper (used internally by every gate):
assert qubit_axis(0, 4) == 3   # qubit 0 is the LSB / rightmost axis
assert qubit_axis(3, 4) == 0   # qubit 3 is the MSB / leftmost axis

# End-to-end Shor.
result = shor_factor(15, seed=42)
print(result)  # ShorFactorResult(p=3, q=5, attempts=1)
```

## Scope and known limitations

- **`shor_factor` does not detect prime powers** (`N = p**k` for prime
  `p`, `k >= 2`). Such inputs make every quantum attempt fail the
  factor-derivation check; the function exhausts `max_attempts` and
  reports failure. The classical pre-check (`round(N**(1/k))**k ==
  N`) belongs in the caller. Even `N` is handled here (short-
  circuits without a quantum step). Matches `/c` and `/go` scope.
- **Shor-25 and beyond are not a v1 performance target.** At 25
  qubits the modexp permutation tensor is 256 MiB (int64) on CPU
  before transfer, and the gather output adds another 512 MiB. The
  arithmetic is correct, but the memory cost is significant.
  Shor-15 (13-qubit register: `n=4` + `t=9`) and Shor-21 (16-qubit
  register: `n=5` + `t=11`) are the supported algorithm targets;
  both run in well under a second on CPU.
- **MPS dtype is `complex64`, not `complex128`.** The MPS backend
  lacks `float64`. `apply_modular_exp`'s permutation tensor and the
  resulting gather output are unaffected (both int64 / complex64),
  but precision-sensitive computations (e.g., very deep gate
  sequences) accumulate rounding faster than on CPU/CUDA.
  `amp_tol_for` and `prob_tol_for` return looser tolerances when
  the register's dtype is `complex64`.
- **No CUDA testing in CI.** GitHub-hosted runners don't have GPUs.
  The device-parametrised tests run CPU-only on CI; MPS coverage
  comes from local runs on Apple Silicon.

## Layout

```
qubit/
  __init__.py            # public API re-exports
  _axis.py               # qubit_axis (LSB-first convention lives here)
  _device.py             # default_device, default_dtype, validate_dtype_device
  _memory.py             # estimate_state_bytes, estimate_peak_bytes, preflight
  _assert.py             # qubit:-prefixed ValueError/TypeError helpers
  _view.py               # state_view + validate_matrix (gate helpers)
  qreg.py                # Qreg class + method wrappers for every gate
  standart.py            # arithmetic helpers (gcd, mod_pow, continued_fraction, ...)
  gates_single.py        # apply_u + apply_h/x/y/z/s/t/phase/rx/ry/rz
  gates_controlled.py    # apply_cu + cnot/cz/controlled_phase/swap
  gates_multi.py         # apply_multi_controlled_z + apply_multi_controlled_x
  measure.py             # measure_qubit / measure_all / sample_distribution / clone / dump
  qft.py                 # apply_qft + apply_qft_inverse (with bit-reversal swaps)
  grover.py              # apply_grover (uniform prep + oracle/diffusion iterations)
  shor.py                # apply_modular_exp + apply_shor_period + shor_factor
  cli.py                 # qubit-demo CLI (argparse, exit codes, top-level catch)
tests/
  conftest.py            # device fixture (parametrises over available devices)
  test_assert.py
  test_axis.py
  test_device.py
  test_memory.py
  test_qreg.py
  test_standart.py
  test_gates_single.py
  test_gates_controlled.py
  test_gates_multi.py
  test_measure.py
  test_qft.py
  test_grover.py
  test_shor.py
  test_cli.py
  test_import.py
pyproject.toml
Makefile
README.md
assessment.md           # thesis-claim coverage map
```

The underscore-prefixed modules are package-private; external callers
should import only from `qubit` (the top-level `__init__.py`
re-exports the public surface).

The misspelled `standart.py` filename is deliberate parity with `/c`'s
2004 spelling and the `/go` sibling. The function names inside use
modern snake_case (`gcd_u64`, `mod_pow`, `continued_fraction`,
`is_power_of_two`, `ilog2_u32`); the `_u64` / `_u32` suffixes are
preserved as visual cues even though Python ints have no width limit.

## Cross-implementation parity

| Claim | This implementation | Sibling locations |
|---|---|---|
| LSB-first basis indexing | `qubit/_axis.py::qubit_axis` | `/go gates_single.go`, `/c parallel.c` |
| Tensor-native gates | `qubit/_view.py::state_view`, `gates_single.apply_u` | `/go parallelOverPairs`, `/c apply_u` |
| ModularExp via permutation | `qubit/shor.py::_build_modexp_perm` | `/go shor.go::_build_modexp_permutation`, `/c shor.c::apply_modular_exp` |
| QFT with bit-reversal swaps | `qubit/qft.py::apply_qft` | `/go qft.go::ApplyQFT`, `/c qft.c::apply_qft` |
| Continued-fraction recovery | `qubit/standart.py::continued_fraction` | `/go standart.go::ContinuedFraction`, `/c standart.c::continued_fraction` |
| Factor N=15 reliably | `tests/test_shor.py::test_shor_factor_15_*` | `/go shor_test.go`, `/c tests/test_shor.c` |
| Shor-21 gated by env var | `tests/test_shor.py::test_shor_period_a2_mod21_gated` | `/go shor_test.go::TestShorPeriodA2Mod21`, `/c tests/test_shor.c` |

See [`assessment.md`](./assessment.md) for the detailed file:line map.

## Sibling implementations

- [/c](../c) — MPI / OpenMPI 5.x. 26-qubit registers tested at NP=1..8.
- [/go](../go) — goroutines / per-call WaitGroup. 70 tests, race-clean.


================================================================================
FILE: implementation/python/assessment.md
================================================================================
# implementation/python — coverage of thesis claims

Updated at the end of Phase 9. File:line references point to the
canonical implementation site for each claim. Sibling implementations:
[`/c`](../c/assessment.md) (MPI) and [`/go`](../go/assessment.md)
(goroutines) cover the same claims via different parallelism models.

## §8 (sparse-gate strategy, 2026 thesis)

| Claim | Status | Location |
|---|---|---|
| Object-level in-place single-qubit gate, O(2^n) work | ✓ | `qubit/gates_single.py::apply_u` |
| State vector as a flat `(2^n,)` PyTorch tensor | ✓ | `qubit/qreg.py::Qreg._amp` |
| Tensor-native dispatch via `tensordot` + `movedim` | ✓ | `qubit/gates_single.py::apply_u` |
| Controlled gates via 4x4 block-diagonal CU | ✓ | `qubit/gates_controlled.py::apply_cu` |
| ModularExp as object-level in-place permutation via `torch.gather` | ✓ | `qubit/shor.py::apply_modular_exp` |
| qreg API per §12 | ✓ | `qubit/qreg.py::Qreg` |

**"In-place" caveat.** Object-level in-place: the `Qreg` instance is
reused across gate calls and the underlying tensor is reassigned. PyTorch
typically allocates a fresh output tensor for `tensordot`, `matmul`, and
`gather`, so transient peak memory is 2x to 4x the state-vector size
during a single gate (see `_memory.estimate_peak_bytes`). The sparse-gate
work is still O(2^n) per gate; we never materialise a 2^n x 2^n operator.

**Parallelism model.** Instead of MPI ranks (`/c`) or goroutines
(`/go`), this implementation routes work through PyTorch's tensor ops
on whatever device is available (NVIDIA CUDA, AMD ROCm, Apple Metal /
MPS, or CPU fallback). The "parallel dispatch" is PyTorch's kernel
scheduler; the simulator never has a `parallelOverPairs`-style
abstraction of its own.

## §9 (2026 QFT)

| Claim | Status | Location |
|---|---|---|
| QFT forward + inverse | ✓ | `qubit/qft.py::apply_qft` / `apply_qft_inverse` |
| Includes final bit-reversal swaps (natural binary order) | ✓ | `qubit/qft.py::apply_qft` final swap loop |
| Inverse is a TRUE inverse, not three forward applications | ✓ | `qubit/qft.py::apply_qft_inverse` (reverses loop order, negates phases) |
| Period detection on known periodic input | ✓ (tested) | `tests/test_qft.py::test_qft_on_basis_state_matches_analytic` |
| Round-trip exhaustively across basis states for n=1..4 | ✓ (tested) | `tests/test_qft.py::test_qft_then_qft_inverse_round_trip` |

## §10 (2026 Grover)

| Claim | Status | Location |
|---|---|---|
| Phase-oracle callback API | ✓ | `qubit/grover.py::apply_grover` (oracle: `Callable[[Qreg, Any], None]`) |
| H^n → oracle / diffusion loop | ✓ | `qubit/grover.py::apply_grover` |
| Diffusion via H X MCZ X H sandwich | ✓ | `qubit/grover.py::apply_grover` |
| Default iteration count = floor(π/4·√N) | ✓ | `qubit/grover.py::apply_grover` |
| Over-iteration drops probability (rotation reverses past π/2) | ✓ (tested) | `tests/test_grover.py::test_grover_over_iteration_drops_probability` |
| Multiple marked items (4 of 16 in 1 iter → P=1.0) | ✓ (tested) | `tests/test_grover.py::test_grover_four_marked_in_16_one_iteration` |
| n_qubits=1 degenerate case supported | ✓ (tested) | `tests/test_grover.py::test_grover_n_qubits_one_does_not_amplify_but_runs` |

## §11 (2026 Shor)

| Claim | Status | Location |
|---|---|---|
| `apply_modular_exp` with y >= N pass-through | ✓ | `qubit/shor.py::apply_modular_exp` |
| `apply_shor_period` (period-finding subroutine) | ✓ | `qubit/shor.py::apply_shor_period` |
| `shor_factor` (end-to-end with retry loop) | ✓ | `qubit/shor.py::shor_factor` |
| Continued-fraction post-processing | ✓ | `qubit/standart.py::continued_fraction` |
| Factor N=15 reliably with seed | ✓ (tested) | `tests/test_shor.py::test_shor_factor_15_seeded` |
| Even N short-circuit (no quantum step) | ✓ (tested) | `tests/test_shor.py::test_shor_factor_even_N_short_circuits` |
| Shor-21 period of 2 mod 21 (gated) | ✓ (gated) | `tests/test_shor.py::test_shor_period_a2_mod21_gated` |
| Reproducibility under seed | ✓ (tested) | `tests/test_shor.py::test_shor_factor_seeded_is_deterministic` |

**Scope caveat.** `shor_factor` does NOT detect prime powers
(`N = p**k` for prime `p`, `k >= 2`). Such inputs make every quantum
attempt fail the factor-derivation check; the function exhausts
`max_attempts` and reports failure. The classical pre-check belongs
in the caller. Matches `/c` and `/go` scope. Documented inline at
`qubit/shor.py::shor_factor` docstring.

## §12 (qreg API)

Every entry in spec §6 is implemented in `qubit/qreg.py` + the
adjacent gate / measurement / algorithm files. Python-specific
deviations from the abstract API:

- **No `Destroy()` / `close()` / context manager.** Python's GC
  reclaims PyTorch tensors when the `Qreg` goes out of scope.
- **Method wrappers for every state-mutating function.** `q.apply_h(0)`
  reads more naturally for method-chaining; `apply_h(q, 0)` matches
  the lower-level function signature. Both are public.
- **Amplitude slice is unexported (`_amp`).** Accessors are
  `amplitude(i)`, `amplitudes_copy()`, `prob_of(basis)`, `norm()`.
- **Functional options at construction.** `Qreg(n, *, device=None,
  seed=None, dtype=None, check_memory=True)`. `None` triggers
  auto-detection per the device/dtype policy.
- **Programmer-error panics use `ValueError` / `TypeError` with the
  `qubit:` prefix.** Same uniform-exception policy across the package;
  no `panic`-vs-`error` split as in `/go`. Library code never calls
  `sys.exit`; the CLI (`qubit/cli.py`) catches at the top level and
  exits non-zero.

## Out of scope for v1

- Distributed execution across machines (would re-introduce MPI's
  problem space; outside the GPU-tensor-ops premise).
- Density matrices / mixed states.
- Noise models.
- Prime-power detection in `shor_factor` (classical pre-check; matches
  `/c` and `/go`).
- Shor-25 and beyond as a routine target (the modexp permutation
  tensor at 25 qubits is 256 MiB on CPU before transfer; the gather
  output is another 512 MiB. The arithmetic is correct, but the
  memory cost is significant. Shor-15 and Shor-21 are the supported
  algorithm targets.)

## Test matrix

`uv run pytest -q` covers **414 active tests** in well under a second on
Apple Silicon (`make check` adds `ruff` and `mypy --strict`). The
gated Shor-21 test brings the total to **415** via `RUN_SHOR_21=1`
and runs the 16-qubit
period-finding circuit with fixed `a=2`; the only stochasticity is
the QFT-readout measurement, so the test asserts the recovered
period divides 6 (the true order of 2 mod 21).

Device coverage on this Apple Silicon host: CPU + MPS via the
parametrised `device` fixture in `tests/conftest.py`. CI runs Linux
CPU-only.

## CLI

The `qubit-demo` console script is exposed via
`[project.scripts] qubit-demo = "qubit.cli:main"` in `pyproject.toml`.
It mirrors the demos in `/c`'s `cmd/qubit` and `/go`'s
`cmd/qubit/main.go`:

```bash
uv run qubit-demo --algo bell
uv run qubit-demo --algo qft
uv run qubit-demo --algo grover
uv run qubit-demo --algo shor
```

All four demos are seeded where applicable so successive runs produce
identical output. Demo errors print to stderr with a `qubit-demo:`
prefix and the process exits non-zero.


================================================================================
FILE: implementation/python/Makefile
================================================================================
# ---------------------------------------------------------------------------
# Makefile for implementation/python -- PyTorch state-vector quantum simulator.
#
# Targets:
#   make sync        uv sync --group dev          (install/refresh .venv)
#   make test        uv run pytest                (full test suite)
#   make test-cpu    PYTORCH_DISABLE_MPS=1 ... (force CPU even on Apple Silicon)
#   make lint        uv run ruff check .
#   make format      uv run ruff format .
#   make typecheck   uv run mypy qubit tests
#   make check       lint + typecheck + test      (the standard PR gate)
#   make demo ALGO=bell|qft|grover|shor
#                    uv run qubit-demo --algo $(ALGO)
#   make clean       rm -rf .venv .pytest_cache .ruff_cache .mypy_cache
#
# The implementation/python module is managed by `uv`; the Makefile is a
# thin convenience wrapper. Most one-off operations can use `uv run <tool>`
# directly.
# ---------------------------------------------------------------------------

ALGO ?= bell

.PHONY: sync test test-cpu lint format typecheck check demo clean

sync:
	uv sync --group dev

test:
	uv run pytest

test-cpu:
	PYTORCH_DISABLE_MPS=1 CUDA_VISIBLE_DEVICES= uv run pytest

lint:
	uv run ruff check .

format:
	uv run ruff format .

typecheck:
	uv run mypy qubit tests

check: lint typecheck test

demo:
	uv run qubit-demo --algo $(ALGO)

clean:
	rm -rf .venv .pytest_cache .ruff_cache .mypy_cache build dist *.egg-info


================================================================================
FILE: implementation/python/pyproject.toml
================================================================================
[project]
name = "qubit"
version = "0.1.0"
description = "PyTorch state-vector quantum simulator (sparse-gate sibling of /c MPI and /go goroutines)"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "torch>=2.4",
]

[project.scripts]
qubit-demo = "qubit.cli:main"

[dependency-groups]
dev = [
    "pytest>=8.0",
    "ruff>=0.6",
    "mypy>=1.11",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["qubit"]

# ---- ruff -----------------------------------------------------------------
[tool.ruff]
target-version = "py312"
line-length = 100

[tool.ruff.lint]
select = [
    "E",   # pycodestyle errors
    "F",   # pyflakes
    "I",   # isort
    "N",   # pep8-naming
    "UP",  # pyupgrade
    "B",   # flake8-bugbear
    "W",   # pycodestyle warnings
]

[tool.ruff.lint.per-file-ignores]
# Package initialisation installs a narrow PyTorch optional-NumPy warning
# filter before importing torch-backed modules.
"qubit/__init__.py" = ["E402"]
# Tests touch the private _amp/_n/_gen fields by design (same-package
# inspection of internal state). Allow underscore-prefix access there.
# Also allow uppercase N in test names and local variables -- Shor's
# algorithm uses uppercase N (the modulus) everywhere in math
# literature, and keeping the convention in tests makes the math-to-
# code traceability obvious.
"tests/*" = ["SLF001", "N802", "N803", "N806"]
# Shor's algorithm uses uppercase `N` (the modulus) universally in
# the math literature; keeping the convention in code makes the
# math-to-code traceability obvious. Same for qreg.py's Shor method
# wrappers which thread `N` through to qubit.shor.
"qubit/shor.py" = ["N803"]
"qubit/qreg.py" = ["N803"]

# ---- mypy -----------------------------------------------------------------
[tool.mypy]
python_version = "3.12"
strict = true
warn_unused_ignores = true
ignore_missing_imports = false

[[tool.mypy.overrides]]
# PyTorch's type stubs are partial; some return types come back as Any.
# Allow that without burning the strict-mode budget elsewhere.
module = "torch.*"
ignore_missing_imports = true
follow_imports = "skip"

[[tool.mypy.overrides]]
# Tests are less strict about Any everywhere.
module = "tests.*"
disallow_untyped_decorators = false
disallow_any_explicit = false

# ---- pytest ---------------------------------------------------------------
[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = ["-ra", "--strict-markers"]


================================================================================
FILE: implementation/python/qubit/__init__.py
================================================================================
"""qubit: PyTorch state-vector quantum-circuit simulator.

Sibling of ``implementation/c`` (MPI) and ``implementation/go``
(goroutines). The PyTorch backend gives device-agnostic GPU support:
the same code runs on NVIDIA CUDA, AMD ROCm, Apple Metal (MPS), and CPU.

Through Phase 8 this ships: the data model (:class:`Qreg` construction,
accessors, the qubit-axis helper), the ``standart`` arithmetic helpers
(gcd, mod_pow, continued_fraction, is_power_of_two, ilog2_u32), the
single-qubit gate primitives (``apply_u`` plus the standard named
gates), the controlled and multi-controlled gates (``apply_cu`` /
``apply_cnot`` / ``apply_cz`` / ``apply_controlled_phase`` /
``apply_swap`` / ``apply_multi_controlled_z`` /
``apply_multi_controlled_x``), the measurement primitives
(``measure_qubit`` / ``measure_all`` / ``sample_distribution`` /
``clone`` / ``dump``), the Quantum Fourier Transform
(``apply_qft`` / ``apply_qft_inverse``), Grover's amplitude
amplification (``apply_grover``), and Shor's algorithm
(``apply_modular_exp`` / ``apply_shor_period`` / ``shor_factor``), and
the ``qubit-demo`` CLI (``qubit.cli``).

Public API surface:

* :class:`Qreg` -- the state-vector register class with methods for
  every gate (``q.apply_h(0)``).
* Function-style gate equivalents (``apply_h(q, 0)``) re-exported from
  :mod:`qubit.gates_single`. Both shapes are first-class.
* :func:`qubit_axis` -- the LSB-first qubit-to-tensor-axis helper.
* Arithmetic helpers from :mod:`qubit.standart`:
  :func:`gcd_u64`, :func:`mod_pow`, :func:`continued_fraction`,
  :func:`is_power_of_two`, :func:`ilog2_u32`.
* Numerical-tolerance constants ``AMP_TOL_C64`` / ``AMP_TOL_C128`` /
  ``PROB_TOL_C64`` / ``PROB_TOL_C128`` plus the dtype-keyed helpers
  ``amp_tol_for`` and ``prob_tol_for`` for test-side use.
"""

from __future__ import annotations

import warnings

# PyTorch warns when its optional NumPy bridge cannot initialize. The
# simulator does not use NumPy, so suppress only this known optional-bridge
# warning while leaving all other PyTorch warnings visible.
warnings.filterwarnings(
    "ignore",
    message=r"Failed to initialize NumPy: No module named 'numpy'.*",
    category=UserWarning,
)

from ._axis import qubit_axis
from .gates_controlled import (
    apply_cnot,
    apply_controlled_phase,
    apply_cu,
    apply_cz,
    apply_swap,
)
from .gates_multi import (
    apply_multi_controlled_x,
    apply_multi_controlled_z,
)
from .gates_single import (
    apply_h,
    apply_phase,
    apply_rx,
    apply_ry,
    apply_rz,
    apply_s,
    apply_t,
    apply_u,
    apply_x,
    apply_y,
    apply_z,
)
from .grover import apply_grover
from .measure import (
    clone,
    dump,
    measure_all,
    measure_qubit,
    sample_distribution,
)
from .qft import (
    apply_qft,
    apply_qft_inverse,
)
from .qreg import (
    AMP_TOL_C64,
    AMP_TOL_C128,
    PROB_TOL_C64,
    PROB_TOL_C128,
    Qreg,
    amp_tol_for,
    prob_tol_for,
)
from .shor import (
    ShorFactorResult,
    ShorPeriodResult,
    apply_modular_exp,
    apply_shor_period,
    shor_factor,
)
from .standart import (
    continued_fraction,
    gcd_u64,
    ilog2_u32,
    is_power_of_two,
    mod_pow,
)

__all__ = [
    "AMP_TOL_C128",
    "AMP_TOL_C64",
    "PROB_TOL_C128",
    "PROB_TOL_C64",
    "Qreg",
    "ShorFactorResult",
    "ShorPeriodResult",
    "amp_tol_for",
    "apply_cnot",
    "apply_controlled_phase",
    "apply_cu",
    "apply_cz",
    "apply_grover",
    "apply_h",
    "apply_modular_exp",
    "apply_multi_controlled_x",
    "apply_multi_controlled_z",
    "apply_phase",
    "apply_qft",
    "apply_qft_inverse",
    "apply_rx",
    "apply_ry",
    "apply_rz",
    "apply_s",
    "apply_shor_period",
    "apply_swap",
    "apply_t",
    "apply_u",
    "apply_x",
    "apply_y",
    "apply_z",
    "clone",
    "continued_fraction",
    "dump",
    "gcd_u64",
    "ilog2_u32",
    "is_power_of_two",
    "measure_all",
    "measure_qubit",
    "mod_pow",
    "prob_tol_for",
    "qubit_axis",
    "sample_distribution",
    "shor_factor",
]


================================================================================
FILE: implementation/python/qubit/_assert.py
================================================================================
"""Programmer-error validation helpers.

Every exception raised by the `qubit` package carries a ``"qubit: "``
prefix so callers can identify simulator-originated errors at a glance
without inspecting tracebacks. Construction-time and gate preconditions
both flow through these helpers; the CLI catches at the top level.
"""

from __future__ import annotations


def raise_value(cond: bool, fmt: str, *args: object) -> None:
    """Raise ``ValueError`` with ``"qubit: " + fmt % args`` if ``cond`` is False.

    Used at the top of every public method for programmer-error
    preconditions (out-of-range qubit indices, control == target, etc.).
    Construction-time input validation (bad ``n_qubits``) goes through the
    same helper. There is no separate ``panic`` vs ``error`` split as in
    /go; Python's exception model is uniform across both layers.
    """
    if not cond:
        raise ValueError("qubit: " + fmt % args)


def raise_type(cond: bool, fmt: str, *args: object) -> None:
    """Raise ``TypeError`` with ``"qubit: " + fmt % args`` if ``cond`` is False.

    Reserved for type-shape failures (caller passed a non-tensor where a
    tensor was expected, etc.). Value-domain failures use
    :func:`raise_value` so the exception class signals which axis of the
    contract was violated.
    """
    if not cond:
        raise TypeError("qubit: " + fmt % args)


================================================================================
FILE: implementation/python/qubit/_axis.py
================================================================================
"""Qubit-to-tensor-axis mapping.

The flat amplitude vector ``q._amp`` is indexed by basis state with the
same convention as ``/c`` and ``/go``: **qubit 0 is the least significant
bit** of the basis index. For a 3-qubit register, the state
``|q2 q1 q0>`` corresponds to ``amp[4*q2 + 2*q1 + q0]``.

When a gate function reshapes the flat vector to the n-D view
``(2,) * n_qubits`` for einsum/tensordot, the tensor axis corresponding
to qubit ``q`` is :func:`qubit_axis` ``= n_qubits - 1 - q`` (so qubit 0
lives on the rightmost / fastest-varying axis). This module is the single
home of that conversion; gates call it instead of recomputing
``n - 1 - q`` ad hoc, so a silent endian flip cannot be introduced one
gate at a time.
"""

from __future__ import annotations

from ._assert import raise_value


def qubit_axis(target: int, n_qubits: int) -> int:
    """Return the tensor-axis index corresponding to qubit ``target``.

    For the ``(2,) * n`` reshape view of a flat ``2**n``-amplitude state
    vector, qubit 0 lives on axis ``n - 1`` (LSB, fastest-varying),
    qubit 1 on axis ``n - 2``, ..., qubit ``n - 1`` on axis 0 (MSB).

    Raises :class:`ValueError` if ``target`` is not in ``[0, n_qubits)``.
    """
    raise_value(
        0 <= target < n_qubits,
        "qubit_axis: target=%d out of [0, %d)",
        target,
        n_qubits,
    )
    return n_qubits - 1 - target


================================================================================
FILE: implementation/python/qubit/_device.py
================================================================================
"""Device selection and dtype policy.

The simulator is device-agnostic: the same source runs on NVIDIA CUDA,
AMD ROCm (uses the CUDA API surface), Apple Metal (MPS), and CPU. The
caller picks a device at :class:`Qreg` construction; this module handles
the default-detection (cuda > mps > cpu) and enforces the one
hardware-driven dtype constraint:

* PyTorch's MPS backend does not support ``float64``, and therefore
  cannot support ``complex128`` (whose real/imag components are
  float64). Calling ``Qreg(..., device='mps', dtype=torch.complex128)``
  must raise rather than silently downgrade. Auto-detect picks the
  right dtype per device when ``dtype=None``.
"""

from __future__ import annotations

import torch

from ._assert import raise_value


def default_device() -> torch.device:
    """Pick the best available device: cuda > mps > cpu."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def default_dtype(device: torch.device) -> torch.dtype:
    """Highest-precision complex dtype the device supports.

    MPS lacks float64, so its complex peak is complex64. Everywhere else
    (CPU, CUDA, ROCm) defaults to complex128 for parity with /c and /go.
    """
    if device.type == "mps":
        return torch.complex64
    return torch.complex128


def validate_dtype_device(device: torch.device, dtype: torch.dtype) -> None:
    """Raise ValueError if (device, dtype) is unsupportable.

    Currently the only banned combination is MPS + complex128. The error
    message tells the caller exactly which two options they have:
    move to CPU for double precision, or accept complex64 on MPS.
    """
    raise_value(
        dtype in (torch.complex64, torch.complex128),
        "validate_dtype_device: dtype=%s must be complex64 or complex128",
        dtype,
    )
    if device.type == "mps" and dtype == torch.complex128:
        raise ValueError(
            "qubit: MPS backend does not support complex128 "
            "(requires float64, which MPS lacks). "
            "Use device='cpu' for double precision, "
            "or dtype=torch.complex64 to stay on MPS."
        )


def coerce_device(device: torch.device | str | None) -> torch.device:
    """Normalise the caller's ``device`` argument to a ``torch.device``.

    ``None`` triggers :func:`default_device`; strings (``"cpu"``,
    ``"cuda:0"``, ``"mps"``) are passed through ``torch.device(...)``.
    A ready-made ``torch.device`` is returned unchanged.
    """
    if device is None:
        return default_device()
    if isinstance(device, str):
        return torch.device(device)
    return device


================================================================================
FILE: implementation/python/qubit/_memory.py
================================================================================
"""Memory preflight helpers.

PyTorch will happily attempt ``torch.zeros(2**40, dtype=complex128)`` and
fail with a low-level runtime error (especially unfriendly on MPS, where
device memory is shared with the host OS). We estimate the peak working
set for a requested operation before allocating and reject obviously-
impossible requests with a clear ``MemoryError``.

Two layers of check:

1. A **sanity ceiling** (1 TiB) applied regardless of device. Catches
   ``n_qubits >= 36`` requests that cannot fit any consumer hardware.
2. A **device-specific free-memory check** when PyTorch exposes one
   (currently CUDA only via ``torch.cuda.mem_get_info``). MPS and CPU
   have no portable free-memory query in stable PyTorch, so the check
   becomes a no-op on those devices and the sanity ceiling is the only
   gate.

Power users can opt out by passing ``check_memory=False`` to
:class:`Qreg`. The preflight is best-effort and not a substitute for
catching ``RuntimeError`` from allocation; it just narrows the cliff.
"""

from __future__ import annotations

from typing import Literal

import torch

# Bytes per amplitude for each supported complex dtype.
_DTYPE_BYTES: dict[torch.dtype, int] = {
    torch.complex64: 8,
    torch.complex128: 16,
}

# Sanity ceiling applied regardless of device. 1 TiB is bigger than any
# consumer GPU (NVIDIA H100 = 80 GiB, M3 Ultra = 192 GiB), so any state
# tensor that exceeds this is unambiguously a bug in the caller's
# parameters rather than a legitimate hardware target.
_SANITY_CEILING: int = 1 << 40  # 1 TiB


# Op tags used by estimate_peak_bytes. State is the baseline (one tensor);
# single_gate / controlled_gate / qft are upper bounds on the transient
# tensors PyTorch allocates during a gate (tensordot / matmul intermediates
# plus an output state vector); modexp allocates a second state tensor and
# an int64 permutation tensor (per spec §5.5).
_Op = Literal["state", "modexp", "qft", "single_gate", "controlled_gate"]


def dtype_bytes(dtype: torch.dtype) -> int:
    """Bytes per amplitude for ``dtype``. Raises on unsupported dtypes."""
    try:
        return _DTYPE_BYTES[dtype]
    except KeyError as exc:
        raise ValueError(
            f"qubit: dtype_bytes: unsupported dtype {dtype} "
            "(expected torch.complex64 or torch.complex128)"
        ) from exc


def estimate_state_bytes(n_qubits: int, dtype: torch.dtype) -> int:
    """Bytes for a single state-vector tensor of ``2**n_qubits`` amplitudes."""
    if n_qubits < 1:
        raise ValueError(
            f"qubit: estimate_state_bytes: n_qubits={n_qubits} must be >= 1"
        )
    return (1 << n_qubits) * dtype_bytes(dtype)


def estimate_peak_bytes(
    n_qubits: int, dtype: torch.dtype, op: _Op = "state"
) -> int:
    """Peak working-set bytes for ``op`` at the given size and dtype.

    * ``"state"`` -- one state vector. Baseline; matches
      :func:`estimate_state_bytes`.
    * ``"single_gate"`` -- ``2 * state``. A single-qubit gate via
      :func:`qubit.gates_single.apply_u` runs ``tensordot`` + ``movedim``,
      which allocates a fresh output state vector alongside the source.
    * ``"controlled_gate"`` -- ``4 * state``. Controlled-U via permute +
      reshape + matmul + reshape + inverse-permute can hold the source,
      the permuted view, a 2x2-block matmul output, and the unpermuted
      result transiently.
    * ``"qft"`` -- ``4 * state``. The QFT decomposes into Hadamards,
      controlled-phase gates, and SWAPs; the controlled-phase steps
      drive the peak. Treated as an upper bound on a single QFT gate;
      a long QFT sequence releases each transient before the next.
    * ``"modexp"`` -- ``2 * state + 2 * int64_permutation``. ModularExp
      allocates a fresh state vector from the gather and holds both a
      CPU-side and a device-side copy of the permutation index tensor
      transiently (the CPU copy is alive until ``.to(device)`` returns).
    """
    state = estimate_state_bytes(n_qubits, dtype)
    if op == "state":
        return state
    if op == "single_gate":
        return 2 * state
    if op == "controlled_gate":
        return 4 * state
    if op == "qft":
        return 4 * state
    if op == "modexp":
        perm = (1 << n_qubits) * 8  # int64
        return 2 * state + 2 * perm
    raise ValueError(f"qubit: estimate_peak_bytes: unknown op={op!r}")


def free_bytes(device: torch.device) -> int | None:
    """Best-effort free memory in bytes for ``device``.

    Returns ``None`` if PyTorch does not expose a free-memory query for
    the device kind. Currently CUDA exposes one via
    ``torch.cuda.mem_get_info``; MPS and CPU do not have a portable
    stable-API query, so this function reports ``None`` for them and the
    sanity ceiling is the only gate.
    """
    if device.type == "cuda":
        free, _total = torch.cuda.mem_get_info(device)
        return int(free)
    return None


def preflight(
    n_qubits: int,
    dtype: torch.dtype,
    device: torch.device,
    op: _Op = "state",
) -> None:
    """Raise ``MemoryError`` if the request can't possibly succeed.

    Applied as two checks: a universal sanity ceiling (1 TiB) and a
    device-specific free-memory check when available. Either failure
    raises with a message that names ``n_qubits``, ``dtype``, ``op``,
    the predicted requirement, and the device.

    No-op when the size is within the sanity ceiling and the device
    free-memory query is unavailable (MPS, CPU).
    """
    need = estimate_peak_bytes(n_qubits, dtype, op)

    if need > _SANITY_CEILING:
        raise MemoryError(
            f"qubit: preflight: requested n_qubits={n_qubits} "
            f"({dtype}, op={op}) needs {need} B but the sanity ceiling "
            f"is {_SANITY_CEILING} B (no consumer hardware fits this; "
            f"pass check_memory=False to bypass)"
        )

    avail = free_bytes(device)
    if avail is None:
        return
    if need > avail:
        raise MemoryError(
            f"qubit: preflight: requested n_qubits={n_qubits} "
            f"({dtype}, op={op}) needs {need} B but device {device} "
            f"has {avail} B free"
        )


================================================================================
FILE: implementation/python/qubit/_view.py
================================================================================
"""Tensor-view helpers for gate primitives.

The simulator stores its state as a flat ``(2**n,)`` tensor at
``Qreg._amp``. Tensor-native gates need to address individual qubits
as axes; the workhorse :func:`state_view` returns the equivalent n-D
view ``(2,) * n`` with qubit ``i`` on axis ``n - 1 - i`` (LSB-first;
see :mod:`qubit._axis`).

Reused by every gate-primitive module. Kept deliberately minimal --
no /go-style "dispatcher" abstraction. PyTorch's tensor ops are
already parallel inside their kernels, so the simulator's "parallel
primitive" is just `torch.tensordot` / `einsum`, not a custom
worker-pool layer.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

if TYPE_CHECKING:
    from .qreg import Qreg


def state_view(q: Qreg) -> torch.Tensor:
    """Return the n-D view of ``q._amp`` with shape ``(2,) * q._n``.

    Qubit ``i`` lives on axis ``n - 1 - i`` (LSB-first, matching
    ``/c``+``/go``). The result is a zero-copy view of the underlying
    contiguous storage; mutating it mutates ``q._amp`` in place. Tensor
    ops like :func:`torch.tensordot` against this view return fresh
    tensors -- they don't alias ``_amp`` -- so gates that want in-place
    updates must reassign ``q._amp`` at the end (see :func:`apply_u`).
    """
    return q._amp.view((2,) * q._n)


def validate_matrix(
    q: Qreg,
    m: torch.Tensor,
    expected_shape: tuple[int, ...],
    fn_name: str,
) -> None:
    """Validate a unitary-matrix argument's shape, dtype, and device.

    Single-qubit gates pass ``expected_shape=(2, 2)``; two-qubit
    controlled gates will pass ``(4, 4)``. The matrix must already be
    a tensor on the same device and dtype as the register -- callers
    are responsible for constructing it with
    ``torch.tensor(..., dtype=q._dtype, device=q._device)``; we don't
    auto-cast or auto-transfer because either would mask caller bugs.

    Raises ``TypeError`` if ``m`` is not a ``torch.Tensor``, otherwise
    ``ValueError`` for any of the three mismatches. All exceptions
    carry the ``"qubit: "`` prefix, with ``fn_name`` as the user-facing
    function tag.
    """
    if not isinstance(m, torch.Tensor):
        raise TypeError(
            f"qubit: {fn_name}: matrix must be a torch.Tensor, "
            f"got {type(m).__name__}"
        )
    if tuple(m.shape) != expected_shape:
        raise ValueError(
            f"qubit: {fn_name}: matrix shape must be {expected_shape}, "
            f"got {tuple(m.shape)}"
        )
    if m.dtype != q._dtype:
        raise ValueError(
            f"qubit: {fn_name}: matrix dtype must be {q._dtype} "
            f"(the register's dtype), got {m.dtype}"
        )
    # Device comparison: torch.device('mps') and torch.device('mps:0')
    # refer to the same physical device but compare unequal under ``==``.
    # Normalise both indices (None means "default index 0" for cuda/mps)
    # before comparing so the matrix-passed-on-device check doesn't
    # falsely reject an obviously-correct caller.
    if m.device.type != q._device.type or (m.device.index or 0) != (
        q._device.index or 0
    ):
        raise ValueError(
            f"qubit: {fn_name}: matrix device must be {q._device} "
            f"(the register's device), got {m.device}"
        )


================================================================================
FILE: implementation/python/qubit/qreg.py
================================================================================
"""Qreg: state-vector quantum register backed by a 1-D PyTorch tensor.

Public surface: construction + accessors, single-qubit gates, controlled
and multi-controlled gates, measurement, QFT, Grover, and Shor. See the
design notes (rounds §1-§5 of the brainstorm) for the overall
architecture and the rationale behind each choice.

The gate methods are thin wrappers around the function-style API in the
``qubit.gates_*``, ``qubit.measure``, ``qubit.qft``, ``qubit.grover``,
and ``qubit.shor`` modules. Both styles are supported deliberately:
``q.apply_h(0)`` reads naturally for method-chaining, while
``apply_h(q, 0)`` works for users who prefer the functional style and
matches the lower-level signature used internally.

Key invariants:

* ``_amp`` is a 1-D contiguous tensor of length ``2**_n``, dtype either
  ``complex64`` (MPS) or ``complex128`` (CPU / CUDA / ROCm).
* The amplitude tensor lives on ``_device``. The Qreg object reference
  is preserved across gate calls, but most gates allocate a fresh
  underlying tensor through PyTorch's ``tensordot`` / ``movedim`` /
  ``matmul`` and reassign ``_amp``. This is "in-place at the Qreg
  object level," not "in-place on the underlying storage": the
  sparse-gate work is still O(2**n), but PyTorch may double the
  transient memory footprint for a single gate. Only
  ``amplitudes_copy()`` intentionally crosses to host.
* A freshly constructed ``Qreg`` is in the computational basis state
  |0...0>; ``amp[0] == 1`` and all other amplitudes are zero. Use
  ``init_basis(b)`` to reset to a different basis state.
* ``_gen`` is a CPU :class:`torch.Generator` regardless of where ``_amp``
  lives. Measurement crosses to CPU at the readout boundary anyway, so
  keeping the RNG on CPU dodges MPS-generator quirks (reports of
  ``Device type MPS is not supported`` and generator-device-mismatch
  errors in the PyTorch tracker) and gives seeded tests bit-identical
  reproducibility across all devices.
* A :class:`Qreg` is **not** safe for concurrent method calls from
  multiple Python threads on the same instance. Different ``Qreg``
  instances are independent.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Final

import torch

from . import _device, _memory
from ._assert import raise_type, raise_value

if TYPE_CHECKING:
    from .shor import ShorPeriodResult

# Numerical tolerances by dtype. Tests pick the appropriate pair from
# the Qreg's dtype property; gate tests will use ``AMP_TOL_*`` to compare
# amplitudes, measurement tests will use ``PROB_TOL_*`` for probabilities.
AMP_TOL_C64: Final[float] = 1e-5
PROB_TOL_C64: Final[float] = 1e-4
AMP_TOL_C128: Final[float] = 1e-10
PROB_TOL_C128: Final[float] = 1e-9


def amp_tol_for(dtype: torch.dtype) -> float:
    """Amplitude-comparison tolerance for the given dtype."""
    if dtype == torch.complex128:
        return AMP_TOL_C128
    if dtype == torch.complex64:
        return AMP_TOL_C64
    raise ValueError(f"qubit: amp_tol_for: unsupported dtype {dtype}")


def prob_tol_for(dtype: torch.dtype) -> float:
    """Probability-comparison tolerance for the given dtype."""
    if dtype == torch.complex128:
        return PROB_TOL_C128
    if dtype == torch.complex64:
        return PROB_TOL_C64
    raise ValueError(f"qubit: prob_tol_for: unsupported dtype {dtype}")


class Qreg:
    """State-vector quantum register.

    Construct with :class:`Qreg(n_qubits) <Qreg>` for the auto-detected
    device and dtype. Keyword arguments override:

    * ``device`` -- one of ``"cpu"``, ``"cuda"``, ``"mps"``, a
      ``torch.device`` instance, or ``None`` (auto-detect).
    * ``seed`` -- integer for the measurement RNG; ``None`` uses fresh
      OS entropy.
    * ``dtype`` -- ``torch.complex64`` or ``torch.complex128``; ``None``
      picks the highest-precision dtype the device supports. ``mps``
      with ``complex128`` raises ``ValueError``.
    * ``check_memory`` -- ``True`` (default) runs a preflight that
      raises ``MemoryError`` for obviously-impossible requests.
      ``False`` skips, leaving raw allocation failures to PyTorch.
    """

    _amp: torch.Tensor       # shape (2**_n,), dtype complex64 or complex128
    _n: int                  # qubit count; >= 1
    _device: torch.device    # device where _amp lives
    _dtype: torch.dtype      # complex64 or complex128
    _gen: torch.Generator    # CPU RNG for measurement sampling

    def __init__(
        self,
        n_qubits: int,
        *,
        device: torch.device | str | None = None,
        seed: int | None = None,
        dtype: torch.dtype | None = None,
        check_memory: bool = True,
    ) -> None:
        # ---- validate caller-facing arguments ---------------------------
        # n_qubits must be a plain int (booleans are ints in Python; reject
        # them explicitly to avoid Qreg(True) silently making a 1-qubit reg).
        raise_type(
            isinstance(n_qubits, int) and not isinstance(n_qubits, bool),
            "Qreg: n_qubits must be int, got %s",
            type(n_qubits).__name__,
        )
        raise_value(
            n_qubits >= 1,
            "Qreg: n_qubits=%d must be >= 1",
            n_qubits,
        )

        # ---- resolve device and dtype -----------------------------------
        resolved_device = _device.coerce_device(device)
        if dtype is None:
            resolved_dtype = _device.default_dtype(resolved_device)
        else:
            # Explicit caller dtype: validate against the device before
            # we commit to allocating.
            _device.validate_dtype_device(resolved_device, dtype)
            resolved_dtype = dtype

        # ---- preflight memory ------------------------------------------
        if check_memory:
            _memory.preflight(
                n_qubits, resolved_dtype, resolved_device, op="state"
            )

        # ---- allocate state vector -------------------------------------
        # zeros() returns a contiguous tensor; later gate code relies on
        # contiguity for the (2,)*n reshape view. Initialise to |0...0>
        # so a freshly constructed Qreg is a valid quantum state (norm 1)
        # and callers can apply gates immediately without an explicit
        # init_basis(0). Use q.init_basis(b) to start from a different
        # basis or to reset between runs.
        amp = torch.zeros(
            1 << n_qubits, dtype=resolved_dtype, device=resolved_device
        )
        amp[0] = 1 + 0j

        # ---- measurement RNG -------------------------------------------
        # CPU generator regardless of where _amp lives. MPS measurement
        # crosses to CPU at the sample step anyway; CUDA measurement
        # likewise pulls the sampled scalar to host before collapsing.
        gen = torch.Generator()
        if seed is None:
            # Default: fresh OS-entropy seed. Different Qregs without a
            # supplied seed produce independent measurement sequences.
            gen.seed()
        else:
            gen.manual_seed(int(seed))

        # ---- commit -----------------------------------------------------
        self._n = n_qubits
        self._device = resolved_device
        self._dtype = resolved_dtype
        self._amp = amp
        self._gen = gen
        # Remembered so later operation-level preflights (e.g. modexp,
        # which transiently allocates a permutation index plus a fresh
        # state-vector copy) can be skipped by callers that already
        # validated memory themselves.
        self._check_memory = check_memory

    # ---- read-only properties --------------------------------------------

    @property
    def n_qubits(self) -> int:
        """Number of qubits this register represents."""
        return self._n

    @property
    def device(self) -> torch.device:
        """Device where the amplitude tensor lives."""
        return self._device

    @property
    def dtype(self) -> torch.dtype:
        """Complex dtype of the amplitude tensor."""
        return self._dtype

    # ---- state-mutating accessor ----------------------------------------

    def init_basis(self, basis: int) -> None:
        """Collapse the register to the computational basis state ``|basis>``.

        Zeros every amplitude, then sets ``amp[basis] = 1 + 0j``. Used to
        prepare a known starting state before applying a gate sequence.
        """
        raise_value(
            0 <= basis < (1 << self._n),
            "init_basis: basis=%d out of [0, %d)",
            basis,
            1 << self._n,
        )
        self._amp.zero_()
        # Tensor assignment of a Python complex to a complex-dtype cell
        # works across CPU / CUDA / MPS. No host->device transfer of a
        # full tensor: this is a single-element write.
        self._amp[basis] = 1 + 0j

    # ---- read-only accessors --------------------------------------------

    def amplitude(self, i: int) -> complex:
        """Return ``amp[i]`` as a Python ``complex``.

        Bounds-checked; raises ``ValueError`` for ``i`` outside
        ``[0, 2**n_qubits)``. The single-element lookup synchronises with
        the device (it has to, to materialise a Python scalar); use this
        for spot checks, not for hot loops. For full-vector inspection
        use :meth:`amplitudes_copy`.
        """
        raise_value(
            0 <= i < (1 << self._n),
            "amplitude: i=%d out of [0, %d)",
            i,
            1 << self._n,
        )
        return complex(self._amp[i].item())

    def amplitudes_copy(self) -> torch.Tensor:
        """Return a detached CPU clone of the full amplitude vector.

        Always returns a CPU tensor regardless of where ``_amp`` lives,
        so callers can inspect, pickle, compare amplitude-by-amplitude
        against ``/c`` and ``/go`` outputs, or feed the result to NumPy
        without further conversion. This intentionally synchronises with
        the device (``.cpu()`` is itself a sync point), so it's a
        diagnostic API, not a hot-path one. Mutations to the returned
        tensor do not affect ``_amp``.
        """
        return self._amp.detach().cpu().clone()

    def prob_of(self, basis: int) -> float:
        """Return ``|amp[basis]|^2`` as a Python ``float``.

        Bounds-checked. Computes ``real**2 + imag**2`` on device (one
        multiply-and-add), then materialises the scalar.
        """
        raise_value(
            0 <= basis < (1 << self._n),
            "prob_of: basis=%d out of [0, %d)",
            basis,
            1 << self._n,
        )
        a = self._amp[basis]
        # Use real()/imag() views and square on device; .item() forces
        # the single-scalar materialisation at the end.
        return float((a.real * a.real + a.imag * a.imag).item())

    def norm(self) -> float:
        """Sum of ``|amp[i]|^2`` over the entire state vector.

        For a normalised state this is 1.0 within floating-point
        precision; gate tests assert ``abs(norm() - 1) <= PROB_TOL``.
        Computed in one ``torch.sum`` call on device.
        """
        amp = self._amp
        return float(
            (amp.real * amp.real + amp.imag * amp.imag).sum().item()
        )

    # ---- single-qubit gate methods --------------------------------------
    #
    # These are thin wrappers around the function-style API in
    # :mod:`qubit.gates_single`. Both call shapes are public:
    # ``q.apply_h(0)`` is method-style; ``apply_h(q, 0)`` (imported from
    # ``qubit``) is the function-style equivalent. They share the same
    # underlying implementation -- the method just forwards self.
    #
    # The :mod:`qubit.gates_single` import is deferred inside each method
    # so importing :mod:`qubit.qreg` from the gates module (for static
    # type hints via ``TYPE_CHECKING``) does not create a runtime cycle.

    def apply_u(self, target: int, u: torch.Tensor) -> None:
        """Apply a 2x2 single-qubit unitary to the target qubit.

        ``u`` must be a 2x2 :class:`torch.Tensor` on the same device and
        with the same complex dtype as this register. See
        :func:`qubit.gates_single.apply_u` for the full contract.
        """
        from . import gates_single

        gates_single.apply_u(self, target, u)

    def apply_h(self, target: int) -> None:
        """Apply the Hadamard gate to ``target``."""
        from . import gates_single

        gates_single.apply_h(self, target)

    def apply_x(self, target: int) -> None:
        """Apply the Pauli-X (bit-flip) gate to ``target``."""
        from . import gates_single

        gates_single.apply_x(self, target)

    def apply_y(self, target: int) -> None:
        """Apply the Pauli-Y gate to ``target``."""
        from . import gates_single

        gates_single.apply_y(self, target)

    def apply_z(self, target: int) -> None:
        """Apply the Pauli-Z (phase-flip) gate to ``target``."""
        from . import gates_single

        gates_single.apply_z(self, target)

    def apply_s(self, target: int) -> None:
        """Apply the S (phase pi/2) gate to ``target``."""
        from . import gates_single

        gates_single.apply_s(self, target)

    def apply_t(self, target: int) -> None:
        """Apply the T (phase pi/4) gate to ``target``."""
        from . import gates_single

        gates_single.apply_t(self, target)

    def apply_phase(self, target: int, theta: float) -> None:
        """Apply the general phase gate ``diag(1, e^{i theta})`` to ``target``."""
        from . import gates_single

        gates_single.apply_phase(self, target, theta)

    def apply_rx(self, target: int, theta: float) -> None:
        """Apply the rotation ``RX(theta)`` to ``target``."""
        from . import gates_single

        gates_single.apply_rx(self, target, theta)

    def apply_ry(self, target: int, theta: float) -> None:
        """Apply the rotation ``RY(theta)`` to ``target``."""
        from . import gates_single

        gates_single.apply_ry(self, target, theta)

    def apply_rz(self, target: int, theta: float) -> None:
        """Apply the rotation ``RZ(theta)`` to ``target``."""
        from . import gates_single

        gates_single.apply_rz(self, target, theta)

    # ---- two-qubit controlled gate methods ------------------------------
    #
    # Same lazy-import discipline as the single-qubit methods. The
    # function-style equivalents live in :mod:`qubit.gates_controlled`.

    def apply_cu(
        self, control: int, target: int, u: torch.Tensor
    ) -> None:
        """Apply a 2x2 controlled-unitary with ``u`` on the target qubit.

        See :func:`qubit.gates_controlled.apply_cu` for the contract.
        """
        from . import gates_controlled

        gates_controlled.apply_cu(self, control, target, u)

    def apply_cnot(self, control: int, target: int) -> None:
        """Apply controlled-NOT: flip ``target`` when ``control`` is 1."""
        from . import gates_controlled

        gates_controlled.apply_cnot(self, control, target)

    def apply_cz(self, control: int, target: int) -> None:
        """Apply controlled-Z: phase-flip when both qubits are 1."""
        from . import gates_controlled

        gates_controlled.apply_cz(self, control, target)

    def apply_controlled_phase(
        self, control: int, target: int, theta: float
    ) -> None:
        """Apply controlled-Phase: multiply ``|11>`` by ``e^{i theta}``."""
        from . import gates_controlled

        gates_controlled.apply_controlled_phase(
            self, control, target, theta
        )

    def apply_swap(self, a: int, b: int) -> None:
        """Exchange qubits ``a`` and ``b`` via the 3-CNOT decomposition."""
        from . import gates_controlled

        gates_controlled.apply_swap(self, a, b)

    # ---- multi-controlled gate methods ----------------------------------

    def apply_multi_controlled_z(
        self, controls: list[int] | tuple[int, ...]
    ) -> None:
        """Phase-flip when every listed control bit is 1."""
        from . import gates_multi

        gates_multi.apply_multi_controlled_z(self, controls)

    def apply_multi_controlled_x(
        self,
        controls: list[int] | tuple[int, ...],
        target: int,
    ) -> None:
        """Flip ``target`` when every listed control is 1 (generalised Toffoli)."""
        from . import gates_multi

        gates_multi.apply_multi_controlled_x(self, controls, target)

    # ---- measurement, sampling, clone, dump -----------------------------

    def measure_qubit(self, target: int) -> int:
        """Projective measurement on ``target``; collapse + renormalise.

        Returns ``0`` or ``1``. See :func:`qubit.measure.measure_qubit`.
        """
        from . import measure

        return measure.measure_qubit(self, target)

    def measure_all(self) -> int:
        """Sample a full basis index from ``|amp|^2`` and collapse.

        See :func:`qubit.measure.measure_all`.
        """
        from . import measure

        return measure.measure_all(self)

    def sample_distribution(self, shots: int) -> list[int]:
        """Run ``shots`` independent measurements; original is unchanged.

        See :func:`qubit.measure.sample_distribution`.
        """
        from . import measure

        return measure.sample_distribution(self, shots)

    def clone(self) -> Qreg:
        """Return an independent copy of this register (amp + RNG state).

        See :func:`qubit.measure.clone`.
        """
        from . import measure

        return measure.clone(self)

    def dump(
        self, *, threshold: float = 0.0
    ) -> list[tuple[int, complex]]:
        """Structured list of ``(basis, amplitude)`` with ``|amp| > threshold``.

        See :func:`qubit.measure.dump`.
        """
        from . import measure

        return measure.dump(self, threshold=threshold)

    # ---- QFT ------------------------------------------------------------

    def apply_qft(self, start: int = 0, n: int | None = None) -> None:
        """Apply the forward QFT to qubits ``[start, start + n)``.

        See :func:`qubit.qft.apply_qft`.
        """
        from . import qft

        qft.apply_qft(self, start, n)

    def apply_qft_inverse(
        self, start: int = 0, n: int | None = None
    ) -> None:
        """Apply the inverse QFT to qubits ``[start, start + n)``.

        See :func:`qubit.qft.apply_qft_inverse`.
        """
        from . import qft

        qft.apply_qft_inverse(self, start, n)

    # ---- Grover ---------------------------------------------------------

    def apply_grover(
        self,
        n_qubits: int,
        oracle: Callable[[Qreg, Any], None],
        user: Any = None,
        iterations: int | None = None,
    ) -> None:
        """Run Grover on qubits ``[0, n_qubits)`` of this register.

        See :func:`qubit.grover.apply_grover`.
        """
        from . import grover

        grover.apply_grover(self, n_qubits, oracle, user, iterations)

    # ---- Shor -----------------------------------------------------------

    def apply_modular_exp(
        self,
        counting_start: int,
        t: int,
        target_start: int,
        n: int,
        a: int,
        N: int,
    ) -> None:
        """Apply the reversible (x, y) -> (x, a^x*y mod N) gate.

        See :func:`qubit.shor.apply_modular_exp`.
        """
        from . import shor

        shor.apply_modular_exp(
            self, counting_start, t, target_start, n, a, N
        )

    def apply_shor_period(
        self,
        counting_start: int,
        t: int,
        target_start: int,
        n: int,
        a: int,
        N: int,
    ) -> ShorPeriodResult:
        """Run one Shor period-finding round.

        See :func:`qubit.shor.apply_shor_period`.
        """
        from . import shor

        return shor.apply_shor_period(
            self, counting_start, t, target_start, n, a, N
        )


================================================================================
FILE: implementation/python/qubit/standart.py
================================================================================
"""Arithmetic helpers shared by gates and algorithms.

The misspelled filename ``standart`` is parity with the original 2004
``/c`` thesis code and the ``/go`` sibling; keeping the misspelling
makes the cross-implementation correspondence visually obvious.

Unlike ``/c`` and ``/go`` -- both of which need ``add_mod``/``mul_mod``
gymnastics to keep ``uint64`` arithmetic from overflowing -- Python's
``int`` is arbitrary-precision, so the operations here are thin wrappers
around stdlib builtins (``math.gcd``, ``pow``,
``fractions.Fraction.limit_denominator``) with a uniform validation
layer that raises ``qubit:``-prefixed exceptions.

Function names keep the ``_u64`` / ``_u32`` suffixes from ``/c``+``/go``
even though Python ints have no such limit; the suffix preserves the
visual parallel and signals the intended non-negative input domain.

Public surface re-exported from :mod:`qubit`:

* :func:`gcd_u64`
* :func:`mod_pow`
* :func:`continued_fraction`
* :func:`is_power_of_two`
* :func:`ilog2_u32`
"""

from __future__ import annotations

import math
from fractions import Fraction

from ._assert import raise_value


def gcd_u64(a: int, b: int) -> int:
    """Return ``gcd(a, b)`` for non-negative ``a`` and ``b``.

    Thin wrapper around :func:`math.gcd` with explicit validation that
    both inputs are non-negative -- ``math.gcd`` accepts negatives but
    the ``/c``+``/go`` siblings work over ``uint64`` and the algorithm
    is only meaningful here for non-negative integers. ``gcd(0, n)``
    returns ``n``; ``gcd(0, 0)`` returns ``0``.
    """
    raise_value(a >= 0, "gcd_u64: a=%d must be non-negative", a)
    raise_value(b >= 0, "gcd_u64: b=%d must be non-negative", b)
    return math.gcd(a, b)


def mod_pow(base: int, exp: int, mod: int) -> int:
    """Return ``base**exp mod mod`` using Python's three-argument :func:`pow`.

    Python's builtin is the canonical fast modular exponentiation; this
    function adds the ``qubit:``-prefixed validation that gates and
    Shor's period finder rely on.

    Validates ``exp >= 0`` and ``mod >= 1``. Negative ``base`` is
    accepted (Python normalises via ``base % mod`` internally), and the
    ``0**0 = 1`` convention from :func:`pow` is preserved.
    """
    raise_value(exp >= 0, "mod_pow: exp=%d must be >= 0", exp)
    raise_value(mod >= 1, "mod_pow: mod=%d must be >= 1", mod)
    return pow(base, exp, mod)


def continued_fraction(x: int, q: int, max_denominator: int) -> int:
    """Return the denominator of the best rational approximation of
    ``x / q`` whose denominator is at most ``max_denominator``.

    Used by Shor's period-finding step: after measuring the counting
    register to get ``c`` in ``[0, 2**t)``, the candidate period ``r``
    is the denominator of the best rational ``p / r`` approximation of
    ``c / 2**t`` with ``r <= N``. This function takes the
    ``(numerator, denominator)`` pair directly instead of a float, so it
    avoids the float-precision tap-dance that ``/c`` and ``/go`` have to
    work around with their early-exit ``frac < 1e-12`` guards.

    Implementation delegates to :meth:`fractions.Fraction.limit_denominator`,
    which performs the same Stern-Brocot convergent search that ``/c``+
    ``/go`` spell out by hand. Validated to match the siblings on every
    Shor case we have (see :file:`tests/test_standart.py`).

    Returns ``1`` for ``x == 0`` (the trivial ``0 / 1`` approximation),
    which matches the ``/go`` ``x <= 0`` safety hatch used at the
    Shor-readout boundary.
    """
    raise_value(q > 0, "continued_fraction: q=%d must be > 0", q)
    raise_value(
        max_denominator > 0,
        "continued_fraction: max_denominator=%d must be > 0",
        max_denominator,
    )
    return Fraction(x, q).limit_denominator(max_denominator).denominator


def is_power_of_two(x: int) -> bool:
    """Return ``True`` iff ``x`` is a positive power of two (1, 2, 4, 8, ...).

    Returns ``False`` for zero and for negative integers. Matches the
    ``/c``+``/go`` ``IsPowerOfTwo`` semantics exactly.
    """
    return x > 0 and (x & (x - 1)) == 0


def ilog2_u32(x: int) -> int:
    """Return ``floor(log2(x))`` for ``x >= 1``.

    Uses :meth:`int.bit_length` minus one -- the canonical Python idiom
    for integer log-base-2. The ``_u32`` suffix preserves parity with
    ``/c``+``/go``, where the argument was constrained to unsigned
    32-bit; Python ints have no such bound, but the name keeps the
    cross-implementation correspondence intact.

    Raises ``ValueError`` for ``x <= 0`` (``log2(0)`` is undefined and
    the function is meaningless on negative inputs).
    """
    raise_value(x > 0, "ilog2_u32: x=%d must be > 0", x)
    return x.bit_length() - 1


================================================================================
FILE: implementation/python/qubit/gates_single.py
================================================================================
"""Single-qubit gate primitives.

Every named gate (``apply_h``, ``apply_x``, ...) builds its 2x2 matrix
on the Qreg's device/dtype and dispatches to :func:`apply_u`, which is
the tensor-native workhorse:

1. ``view = q._amp.view((2,) * n)`` -- zero-copy n-D reshape.
2. ``out = torch.tensordot(u, view, dims=([1], [axis]))`` -- contract
   ``u``'s input axis with the target qubit's tensor axis.
3. ``out.movedim(0, axis)`` -- move ``u``'s output axis back to where
   the target was, restoring the original layout.
4. ``q._amp = out.reshape(-1)`` -- flatten back to ``(2**n,)``. This
   step **allocates a fresh state vector** because ``movedim``'s
   strides are non-row-major; ``reshape(-1)`` materialises a contiguous
   copy. One state-vector worth of allocation per gate.

The LSB-first axis mapping lives in :func:`qubit._axis.qubit_axis`. The
matrix-on-correct-device-and-dtype check is :func:`qubit._view.validate_matrix`.
Both are reused by the controlled and multi-controlled gate modules in
later phases.
"""

from __future__ import annotations

import cmath
import math
from typing import TYPE_CHECKING

import torch

from ._axis import qubit_axis
from ._view import state_view, validate_matrix

if TYPE_CHECKING:
    from .qreg import Qreg


# ---------------------------------------------------------------------------
# apply_u: the workhorse every other gate routes through.
# ---------------------------------------------------------------------------


def apply_u(q: Qreg, target: int, u: torch.Tensor) -> None:
    """Apply the 2x2 single-qubit unitary ``u`` to ``q``'s ``target`` qubit.

    ``u`` must be a 2x2 ``torch.Tensor`` on the same device and with the
    same complex dtype as ``q._amp``. Caller is responsible for
    constructing it; see the named-gate functions below for examples.

    Modifies ``q`` in place: ``q._amp`` is replaced with a freshly
    allocated tensor holding the result. See module docstring for the
    allocation accounting.
    """
    # qubit_axis validates target range and raises a qubit:-prefixed
    # ValueError on out-of-range. matrix-validation is via _view.
    axis = qubit_axis(target, q._n)
    validate_matrix(q, u, (2, 2), "apply_u")

    # Zero-copy n-D view; the underlying storage is q._amp's flat buffer.
    view = state_view(q)

    # tensordot semantics: contract u's axis 1 (column / "input qubit
    # value") with view's `axis` (the target's qubit axis). Result has
    # u's axis 0 (row / "output qubit value") prepended, followed by
    # view's remaining axes in their original order with `axis` removed.
    out = torch.tensordot(u, view, dims=([1], [axis]))

    # The new target-qubit axis is at position 0 after tensordot; move
    # it back to where the original target axis was so the result has
    # the original layout. (movedim returns a non-contiguous view.)
    out = out.movedim(0, axis)

    # Materialise as a fresh contiguous (2^n,) tensor. movedim's
    # non-row-major strides force reshape(-1) to allocate; the old
    # q._amp becomes garbage and is reclaimed on the next GC cycle.
    q._amp = out.reshape(-1)


# ---------------------------------------------------------------------------
# Pauli and Hadamard.
# ---------------------------------------------------------------------------


def apply_h(q: Qreg, target: int) -> None:
    """Apply the Hadamard gate ``H = (1/sqrt(2)) [[1, 1], [1, -1]]``."""
    inv_sqrt2 = 1.0 / math.sqrt(2.0)
    h = torch.tensor(
        [[inv_sqrt2 + 0j, inv_sqrt2 + 0j],
         [inv_sqrt2 + 0j, -inv_sqrt2 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, h)


def apply_x(q: Qreg, target: int) -> None:
    """Apply the Pauli-X (bit-flip) gate ``X = [[0, 1], [1, 0]]``."""
    x = torch.tensor(
        [[0 + 0j, 1 + 0j],
         [1 + 0j, 0 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, x)


def apply_y(q: Qreg, target: int) -> None:
    """Apply the Pauli-Y gate ``Y = [[0, -i], [i, 0]]``."""
    y = torch.tensor(
        [[0 + 0j, -1j],
         [1j, 0 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, y)


def apply_z(q: Qreg, target: int) -> None:
    """Apply the Pauli-Z (phase-flip) gate ``Z = [[1, 0], [0, -1]]``."""
    z = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, -1 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, z)


# ---------------------------------------------------------------------------
# Phase gates: S = phase(pi/2), T = phase(pi/4), general phase(theta).
# ---------------------------------------------------------------------------


def apply_s(q: Qreg, target: int) -> None:
    """Apply the S gate (phase pi/2): ``S = [[1, 0], [0, i]]``.

    Satisfies ``S**2 = Z``.
    """
    s = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, 1j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, s)


def apply_t(q: Qreg, target: int) -> None:
    """Apply the T gate (phase pi/4): ``T = [[1, 0], [0, e^{i pi/4}]]``.

    Satisfies ``T**4 = Z`` and ``T**2 = S``.
    """
    phase_factor = cmath.exp(complex(0, math.pi / 4))
    t = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, phase_factor]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, t)


def apply_phase(q: Qreg, target: int, theta: float) -> None:
    """Apply the general phase gate ``Phase(theta) = diag(1, e^{i theta})``."""
    phase_factor = cmath.exp(complex(0, theta))
    p = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, phase_factor]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, p)


# ---------------------------------------------------------------------------
# Rotations: rx / ry / rz around the Bloch sphere axes.
# ---------------------------------------------------------------------------


def apply_rx(q: Qreg, target: int, theta: float) -> None:
    """Apply ``RX(theta) = [[cos(t/2), -i sin(t/2)], [-i sin(t/2), cos(t/2)]]``.

    ``RX(2 pi)`` maps ``|0>`` to ``-|0>`` (the well-known sign flip after
    a full rotation).
    """
    c = math.cos(theta / 2.0)
    s = math.sin(theta / 2.0)
    off = complex(0, -s)  # -i * sin(theta/2)
    rx = torch.tensor(
        [[c + 0j, off],
         [off, c + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, rx)


def apply_ry(q: Qreg, target: int, theta: float) -> None:
    """Apply ``RY(theta) = [[cos(t/2), -sin(t/2)], [sin(t/2), cos(t/2)]]``.

    Real-valued; satisfies ``RY(4 pi) = I``.
    """
    c = math.cos(theta / 2.0)
    s = math.sin(theta / 2.0)
    ry = torch.tensor(
        [[c + 0j, -s + 0j],
         [s + 0j, c + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, ry)


def apply_rz(q: Qreg, target: int, theta: float) -> None:
    """Apply ``RZ(theta) = diag(e^{-i theta/2}, e^{i theta/2})``.

    Diagonal in the computational basis. ``RZ(theta) |0> = e^{-i theta/2} |0>``.
    """
    neg_half = cmath.exp(complex(0, -theta / 2.0))
    pos_half = cmath.exp(complex(0, theta / 2.0))
    rz = torch.tensor(
        [[neg_half, 0 + 0j],
         [0 + 0j, pos_half]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_u(q, target, rz)


================================================================================
FILE: implementation/python/qubit/gates_controlled.py
================================================================================
"""Two-qubit controlled gates.

The workhorse is :func:`apply_cu`: given a 2x2 ``u`` and ``(control,
target)`` qubits, it embeds ``u`` into a 4x4 controlled unitary
``CU = diag(I, U)`` and applies it to the (control, target) sub-pair
of the state vector.

Convention: the 4x4 acts on the two-qubit subspace in the basis order

    |control, target> = |00>, |01>, |10>, |11>

so the ``|control = 1>`` block sits at rows/cols ``2:4`` (where ``u``
goes) and the ``|control = 0>`` block at rows/cols ``0:2`` is the
identity.

Implementation strategy (tensor-native, no per-amplitude loops):

1. View the flat state as ``(2,) * n``.
2. Permute axes so the control axis lands at position 0 and the target
   axis at position 1, with remaining axes in their original order.
3. Reshape to ``(4, rest)`` and matrix-multiply ``cu @ front``.
4. Reshape back to ``(2, 2) + (2,) * (n - 2)`` and apply the inverse
   permutation to restore the original axis layout.
5. Replace ``q._amp`` with a fresh contiguous ``(2**n,)`` tensor.

The two reshapes are allocation points (post-permute non-row-major
strides force ``reshape`` to materialise a contiguous copy). One
state-vector worth of allocation per controlled-gate call.
"""

from __future__ import annotations

import cmath
from typing import TYPE_CHECKING

import torch

from ._assert import raise_value
from ._axis import qubit_axis
from ._view import state_view, validate_matrix

if TYPE_CHECKING:
    from .qreg import Qreg


def _bring_to_front(n: int, axes: list[int]) -> tuple[list[int], list[int]]:
    """Return (perm, inv_perm) that bring ``axes`` to positions 0..len(axes)-1.

    Remaining axes follow in their original order. The inverse permutation
    is built so ``permute(perm).permute(inv_perm)`` is the identity.
    """
    remaining = [a for a in range(n) if a not in axes]
    perm = list(axes) + remaining
    inv_perm = [0] * n
    for new_idx, old_axis in enumerate(perm):
        inv_perm[old_axis] = new_idx
    return perm, inv_perm


# ---------------------------------------------------------------------------
# apply_cu: the controlled-U workhorse.
# ---------------------------------------------------------------------------


def apply_cu(
    q: Qreg, control: int, target: int, u: torch.Tensor
) -> None:
    """Apply the controlled-U gate with ``u`` (2x2) on ``target`` when ``control`` is 1.

    ``u`` must be a 2x2 ``torch.Tensor`` on the register's device with the
    register's dtype. See :mod:`qubit.gates_single.apply_u` for the rationale.

    Internally builds the 4x4 controlled unitary ``CU = diag(I_2, u)`` and
    applies it to the (control, target) two-qubit subspace via a permute +
    reshape + matmul + reshape + inverse-permute sequence; no per-amplitude
    Python loops.
    """
    raise_value(
        0 <= control < q._n,
        "apply_cu: control=%d out of [0, %d)", control, q._n,
    )
    raise_value(
        0 <= target < q._n,
        "apply_cu: target=%d out of [0, %d)", target, q._n,
    )
    raise_value(
        control != target,
        "apply_cu: control == target == %d", control,
    )
    validate_matrix(q, u, (2, 2), "apply_cu")

    n = q._n

    # Build CU on the same device/dtype. Basis order is |control, target>:
    # |00>, |01>, |10>, |11>. The |control=1> block (rows/cols 2:4) holds u.
    cu = torch.eye(4, dtype=q._dtype, device=q._device)
    cu[2:4, 2:4] = u

    a_control = qubit_axis(control, n)
    a_target = qubit_axis(target, n)
    perm, inv_perm = _bring_to_front(n, [a_control, a_target])

    view = state_view(q)

    # Permute -> reshape to (4, rest). reshape allocates because the
    # permute output is non-contiguous; the resulting tensor is
    # contiguous so the subsequent matmul output strides are clean.
    front = view.permute(*perm).reshape(4, -1)
    front = cu @ front

    # Back to (2, 2) + remaining-axis shape, then inverse permute to the
    # original axis ordering, then flatten.
    out = front.reshape((2, 2) + (2,) * (n - 2))
    out = out.permute(*inv_perm)
    q._amp = out.reshape(-1)


# ---------------------------------------------------------------------------
# Named controlled gates.
# ---------------------------------------------------------------------------


def apply_cnot(q: Qreg, control: int, target: int) -> None:
    """Apply controlled-NOT (CX): flip ``target`` when ``control`` is 1."""
    x = torch.tensor(
        [[0 + 0j, 1 + 0j],
         [1 + 0j, 0 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_cu(q, control, target, x)


def apply_cz(q: Qreg, control: int, target: int) -> None:
    """Apply controlled-Z: phase-flip when both qubits are 1."""
    z = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, -1 + 0j]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_cu(q, control, target, z)


def apply_controlled_phase(
    q: Qreg, control: int, target: int, theta: float
) -> None:
    """Apply controlled-Phase: multiply the ``|11>`` branch by ``e^{i theta}``.

    Specialisations: ``theta = pi`` recovers controlled-Z; ``theta = pi/2``
    is the controlled-S used by QFT.
    """
    phase_factor = cmath.exp(complex(0, theta))
    p = torch.tensor(
        [[1 + 0j, 0 + 0j],
         [0 + 0j, phase_factor]],
        dtype=q._dtype,
        device=q._device,
    )
    apply_cu(q, control, target, p)


def apply_swap(q: Qreg, a: int, b: int) -> None:
    """Exchange the states of qubits ``a`` and ``b`` via the 3-CNOT identity.

    ``SWAP = CNOT(a, b) -> CNOT(b, a) -> CNOT(a, b)``. Same decomposition
    as ``/c`` and ``/go``. Raises if ``a == b`` (caller bug) -- a no-op
    would silently hide the issue.
    """
    raise_value(
        0 <= a < q._n, "apply_swap: a=%d out of [0, %d)", a, q._n
    )
    raise_value(
        0 <= b < q._n, "apply_swap: b=%d out of [0, %d)", b, q._n
    )
    raise_value(a != b, "apply_swap: a == b == %d", a)
    apply_cnot(q, a, b)
    apply_cnot(q, b, a)
    apply_cnot(q, a, b)


================================================================================
FILE: implementation/python/qubit/gates_multi.py
================================================================================
"""Multi-controlled gates.

Two operations: :func:`apply_multi_controlled_z` (phase-flip when every
listed control bit is 1; used by Grover's diffusion) and
:func:`apply_multi_controlled_x` (generalised Toffoli; flip ``target``
when every control bit is 1).

Both are expressed as vectorised mask/gather/scatter ops on the flat
``q._amp`` rather than as n-D tensordot. Reason: the operation only
touches a small subset of amplitudes (the ones whose control bits are
all 1), so building a boolean index over the full ``2**n`` indices and
using PyTorch's fancy indexing is both clearer and lighter on memory
than the permute-and-block-multiply pattern used for two-qubit gates.

Bitwise integer ops have inconsistent backend coverage (CUDA / ROCm /
MPS), so the bitmask is built on CPU and ``.to(device)``-d before the
final indexing step. This is cheap -- the mask is int64, half the
size of complex128 and the same size as complex64.

Multi-controlled X (Toffoli) gotcha: the mask must select **only**
amplitudes where every control is 1 **and** the target bit is 0. Each
such index is then paired with its target-bit-flipped sibling, and the
two values are swapped. Selecting both target=0 and target=1 states
would either double-swap each pair (no-op) or collide depending on
order. Discussed inline at :func:`apply_multi_controlled_x`.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from ._assert import raise_value

if TYPE_CHECKING:
    from .qreg import Qreg


def _validate_controls(
    controls: list[int] | tuple[int, ...],
    n_qubits: int,
    fn_name: str,
    target: int | None = None,
) -> int:
    """Validate the controls list and return the OR-mask of their bit positions.

    Catches: empty list, out-of-range index, duplicates, collision with
    ``target`` (if supplied). All raise ``ValueError`` with the
    ``"qubit: <fn_name>: ..."`` prefix.
    """
    raise_value(
        len(controls) > 0,
        "%s: controls must be a non-empty sequence", fn_name,
    )
    seen: set[int] = set()
    c_mask = 0
    for c in controls:
        raise_value(
            0 <= c < n_qubits,
            "%s: control=%d out of [0, %d)", fn_name, c, n_qubits,
        )
        raise_value(
            c not in seen,
            "%s: duplicate control=%d in %r", fn_name, c, list(controls),
        )
        if target is not None:
            raise_value(
                c != target,
                "%s: control %d == target", fn_name, c,
            )
        seen.add(c)
        c_mask |= 1 << c
    return c_mask


# ---------------------------------------------------------------------------
# Multi-controlled Z: phase-flip when every control bit is 1.
# ---------------------------------------------------------------------------


def apply_multi_controlled_z(
    q: Qreg, controls: list[int] | tuple[int, ...]
) -> None:
    """Phase-flip every amplitude whose listed control bits are all 1.

    The standard use case is Grover's diffusion step, where ``controls``
    is ``list(range(n))`` over the searched qubits. The list-based
    signature lets callers mark arbitrary subsets, which is the more
    general primitive even though Grover only needs the contiguous case.

    Raises ``ValueError`` for an empty ``controls`` sequence, an
    out-of-range index, or duplicate control indices.
    """
    c_mask = _validate_controls(
        controls, q._n, "apply_multi_controlled_z"
    )

    # Build the boolean mask on CPU -- bitwise integer ops on accelerator
    # backends are not uniformly available -- and transfer to the
    # register's device. The mask covers every "controls-all-1" basis
    # index; everything else is left alone.
    idx = torch.arange(1 << q._n, dtype=torch.int64)
    mask_cpu = (idx & c_mask) == c_mask
    mask = mask_cpu.to(q._device)

    # Negate the selected amplitudes. PyTorch's masked write supports
    # complex dtype on every backend we target.
    q._amp[mask] = -q._amp[mask]


# ---------------------------------------------------------------------------
# Multi-controlled X (generalised Toffoli): flip target when all controls = 1.
# ---------------------------------------------------------------------------


def apply_multi_controlled_x(
    q: Qreg,
    controls: list[int] | tuple[int, ...],
    target: int,
) -> None:
    """Flip ``target`` when every listed control is 1 (generalised Toffoli).

    With ``len(controls) == 1`` this is exactly :func:`apply_cnot`; with
    ``len(controls) == 2`` it's the standard Toffoli; higher arities
    generalise.

    Raises ``ValueError`` for an out-of-range ``target``, empty
    ``controls``, out-of-range control, duplicate control, or any control
    equal to ``target``.

    Implementation note: the mask is built to select **only** amplitudes
    with every control bit 1 AND the target bit 0. Each such index is
    paired with its target-bit-flipped sibling (``idx | t_bit``), and the
    two amplitudes are swapped. Selecting both sides of every pair would
    double-swap (no-op) or collide depending on PyTorch's evaluation
    order; selecting only one side is the correct generalisation of
    ``/c`` and ``/go``'s pair-iteration trick.
    """
    raise_value(
        0 <= target < q._n,
        "apply_multi_controlled_x: target=%d out of [0, %d)",
        target, q._n,
    )
    c_mask = _validate_controls(
        controls, q._n, "apply_multi_controlled_x", target=target
    )
    t_bit = 1 << target

    # Build indices on CPU (portability), then move to the device.
    idx = torch.arange(1 << q._n, dtype=torch.int64)
    # Pair selector: every "controls-all-1 AND target-0" index. Each one
    # pairs with `index | t_bit` (the "controls-all-1 AND target-1"
    # sibling), and we swap their amplitudes.
    swap_from_mask = ((idx & c_mask) == c_mask) & ((idx & t_bit) == 0)
    swap_from = idx[swap_from_mask].to(q._device)
    swap_to = (idx[swap_from_mask] | t_bit).to(q._device)

    # Fancy indexing on read returns a fresh tensor (no aliasing with
    # q._amp), so the saved "from-side" amplitudes survive the
    # to-side -> from-side write that comes next.
    from_values = q._amp[swap_from].clone()
    q._amp[swap_from] = q._amp[swap_to]
    q._amp[swap_to] = from_values


================================================================================
FILE: implementation/python/qubit/measure.py
================================================================================
"""Measurement, sampling, clone, and diagnostic dump.

Five public functions, all also exposed as :class:`Qreg` methods:

* :func:`measure_qubit` -- single-qubit projective measurement; the
  register collapses onto the measured branch and renormalises.
* :func:`measure_all` -- sample a full basis index from the
  ``|amp|^2`` distribution and collapse onto it.
* :func:`sample_distribution` -- run ``shots`` independent measurements
  *without* mutating the original register (snapshot + restore).
* :func:`clone` -- independent copy with the same amplitudes and RNG
  state. Mutations on the clone don't reach the original; both
  generators produce the same future sequence until one advances.
* :func:`dump` -- structured diagnostic: list of
  ``(basis_index, amplitude)`` pairs with ``|amp| > threshold``.

Measurement is the documented "device -> host" boundary: we deliberately
``.cpu()`` the probability vector once per call to feed the CPU
:class:`torch.Generator` (the cross-device-portable RNG choice from
Phase 0+1). Gate ops stay on device; only the measurement readout
crosses to host.
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

import torch

from ._assert import raise_value
from ._axis import qubit_axis

if TYPE_CHECKING:
    from .qreg import Qreg


# ---------------------------------------------------------------------------
# measure_qubit: single-qubit projective measurement with collapse.
# ---------------------------------------------------------------------------


def measure_qubit(q: Qreg, target: int) -> int:
    """Perform a projective measurement on the ``target`` qubit.

    Returns ``0`` or ``1``. The register is collapsed onto the chosen
    branch and renormalised so :meth:`Qreg.norm` remains 1.

    Uses the per-Qreg CPU RNG ``q._gen`` for the sample. Two scalar
    syncs to host: the two slice-probability sums (so we can sample on
    CPU). Raises :class:`RuntimeError` with the ``qubit:`` prefix if
    the chosen branch has probability ``<= 0``, which only happens if
    the state has drifted from normalisation.
    """
    raise_value(
        0 <= target < q._n,
        "measure_qubit: target=%d out of [0, %d)", target, q._n,
    )

    n = q._n
    axis = qubit_axis(target, n)

    # n-D view; select gives non-contiguous views of the two halves
    # along the target axis. In-place ops on these slices modify
    # q._amp's underlying storage.
    view = q._amp.view((2,) * n)
    slice_0 = view.select(axis, 0)
    slice_1 = view.select(axis, 1)

    # Compute both probabilities from the actual slice contents (rather
    # than p1 = 1 - p0) so the renormalisation is exact relative to the
    # chosen branch even if the state has drifted slightly.
    p0 = float(
        (slice_0.real * slice_0.real + slice_0.imag * slice_0.imag).sum().item()
    )
    p1 = float(
        (slice_1.real * slice_1.real + slice_1.imag * slice_1.imag).sum().item()
    )

    # Guard zero-norm states (e.g. a fresh Qreg before init_basis) up
    # front, so the sampling step below cannot pick a zero-probability
    # branch and divide by sqrt(0) at the renormalisation step.
    total = p0 + p1
    if total <= 0:
        raise RuntimeError(
            f"qubit: measure_qubit: total probability {total:.6g} <= 0; "
            f"the state has not been initialised or has drifted from "
            f"normalisation (check q.norm())"
        )

    # CPU random draw; q._gen is the per-Qreg CPU generator from Phase 0+1.
    # Sample against the normalised CDF p0/total so an artificially scaled
    # (or slightly drifted) state still produces an unbiased outcome:
    # plain `u < p0` would silently misbehave for any state where
    # p0 + p1 != 1 (the obvious case is a scaled state with p0 = p1 = 2
    # where the un-normalised comparison u < 2 is always True).
    u = float(torch.rand((), generator=q._gen).item())
    outcome = 0 if u < (p0 / total) else 1
    p_outcome = p0 if outcome == 0 else p1

    # Given total > 0 above, the chosen-branch probability is itself > 0
    # because `u < p0/total` cannot be True when p0 == 0 (the CDF cut
    # would be 0 and u is in [0, 1)), and likewise outcome cannot be 1
    # when p1 == 0 (the CDF cut would be 1.0 and u < 1.0 is always True).
    # The check below is defence in depth for floating-point oddities.
    if p_outcome <= 0:
        raise RuntimeError(
            f"qubit: measure_qubit: outcome={outcome} sampled with "
            f"p_outcome={p_outcome:.6g} <= 0; this should not happen "
            f"given total={total:.6g} > 0 (numerical edge case)"
        )

    inv_norm = 1.0 / math.sqrt(p_outcome)

    # Collapse: zero the other half, rescale the chosen half. Both
    # operations are in-place on slices that view q._amp's storage,
    # so q._amp is updated without any reshape/flatten work.
    if outcome == 0:
        slice_1.zero_()
        slice_0.mul_(inv_norm)
    else:
        slice_0.zero_()
        slice_1.mul_(inv_norm)

    return outcome


# ---------------------------------------------------------------------------
# measure_all: sample a full basis index, collapse to it.
# ---------------------------------------------------------------------------


def measure_all(q: Qreg) -> int:
    """Sample a full basis index from the ``|amp|^2`` distribution.

    The register collapses to a pure basis state at the chosen index;
    after this call ``q.norm()`` is 1.0 and exactly one amplitude is
    ``1 + 0j``.

    The probability vector is moved to CPU once so :func:`torch.multinomial`
    can use the per-Qreg CPU generator.
    """
    amp = q._amp
    # Real-valued probability vector on device, then one .cpu() sync
    # so multinomial can pull a sample using q._gen (CPU generator).
    # multinomial accepts non-normalised weights -- it renormalises
    # internally -- so we don't pre-normalise here, but we DO catch
    # the zero-norm case before calling multinomial: otherwise
    # torch.multinomial raises a raw RuntimeError without the
    # "qubit: " prefix, leaking a generic torch error to callers.
    probs = (amp.real * amp.real + amp.imag * amp.imag).detach().cpu()
    total = float(probs.sum().item())
    if total <= 0:
        raise RuntimeError(
            f"qubit: measure_all: total probability {total:.6g} <= 0; "
            f"the state has not been initialised or has drifted from "
            f"normalisation (check q.norm())"
        )

    chosen = int(torch.multinomial(probs, 1, generator=q._gen).item())

    # Collapse on device: zero everything, set the chosen index to 1+0j.
    q._amp.zero_()
    q._amp[chosen] = 1 + 0j
    return chosen


# ---------------------------------------------------------------------------
# sample_distribution: shots independent measurements, original untouched.
# ---------------------------------------------------------------------------


def sample_distribution(q: Qreg, shots: int) -> list[int]:
    """Draw ``shots`` samples from ``|amp|^2`` and return them as a list.

    Does **not** mutate ``q``: the amplitude tensor is read once into a
    probability vector and ``torch.multinomial`` produces all samples
    from it in one call, so the register is never collapsed. ``q._gen``
    is advanced by exactly ``shots`` draws.

    O(2**n + shots) work and O(2**n) extra memory for the probability
    vector. (The previous implementation cloned the full amplitude
    tensor once per shot, which scaled as O(shots * 2**n) instead.)

    ``shots=0`` is valid and returns the empty list. ``shots < 0`` raises.
    Returns a plain ``list[int]`` for ergonomic use with
    :class:`collections.Counter`.
    """
    raise_value(
        shots >= 0,
        "sample_distribution: shots=%d must be >= 0", shots,
    )
    if shots == 0:
        return []

    # Build the probability vector once. q._gen lives on CPU so we move
    # the probs over too -- multinomial requires both on the same device
    # as the generator. For complex64/128 the cast to real-valued probs
    # is straightforward.
    amp = q._amp.detach()
    probs = (amp.real * amp.real + amp.imag * amp.imag).to(
        device="cpu", dtype=torch.float64
    )
    total = float(probs.sum().item())
    if total <= 0.0:
        raise RuntimeError(
            "qubit: sample_distribution: total probability is zero "
            "(forgot to init_basis(), or non-unitary gate sequence?)"
        )

    samples = torch.multinomial(
        probs, num_samples=shots, replacement=True, generator=q._gen
    )
    return [int(x) for x in samples.tolist()]


# ---------------------------------------------------------------------------
# clone: independent Qreg with same amplitudes and RNG state.
# ---------------------------------------------------------------------------


def clone(q: Qreg) -> Qreg:
    """Return an independent copy of ``q``.

    The clone has its own ``_amp`` tensor (a detached copy of ``q._amp``)
    and its own ``_gen`` generator initialised with the same state as
    ``q._gen``. From the moment of cloning, the two registers are
    independent: mutating one's amplitudes doesn't affect the other's,
    and advancing one's RNG (by measurement) doesn't advance the
    other's. They produce the **same** future random sequence until
    one of them draws.

    Bypasses :class:`Qreg`'s ``__init__`` (via ``__new__``) to avoid
    the otherwise-wasted ``torch.zeros(2**n)`` allocation that would
    immediately get overwritten by the amplitude copy. Same trick is
    used by :func:`sample_distribution` indirectly via its snapshot
    pattern.
    """
    from .qreg import Qreg as _Qreg

    new_q = _Qreg.__new__(_Qreg)
    new_q._n = q._n
    new_q._device = q._device
    new_q._dtype = q._dtype
    new_q._amp = q._amp.detach().clone()

    # Match RNG state. Both generators are independent objects from
    # here on; only their initial state is shared.
    new_q._gen = torch.Generator()
    new_q._gen.set_state(q._gen.get_state())
    return new_q


# ---------------------------------------------------------------------------
# dump: structured diagnostic list.
# ---------------------------------------------------------------------------


def dump(q: Qreg, *, threshold: float = 0.0) -> list[tuple[int, complex]]:
    """Return a list of ``(basis_index, amplitude)`` pairs with ``|amp| > threshold``.

    Diagnostic API; does no printing (callers can pretty-print the list
    however they like). Threshold defaults to 0, which returns every
    non-zero amplitude. Pass a positive threshold to filter out small
    amplitudes that you don't care about.

    Pulls the amplitude tensor to CPU once via ``amplitudes_copy``-style
    detach/copy. The boolean mask + ``nonzero`` scan still costs
    ``O(2**n)`` work to find which amplitudes to emit; only the final
    Python ``list`` construction is bounded by the number of selected
    amplitudes. Treat ``dump()`` as a small-state diagnostic helper,
    not a sparse iterator for large registers.
    """
    raise_value(
        threshold >= 0,
        "dump: threshold=%g must be non-negative", threshold,
    )

    cpu_amp = q._amp.detach().cpu()
    # Find indices whose |amp| exceeds the threshold. .abs() on a
    # complex tensor returns magnitudes; the boolean mask is on the
    # real magnitude space.
    mask = cpu_amp.abs() > threshold
    indices = mask.nonzero(as_tuple=True)[0].tolist()
    return [(int(i), complex(cpu_amp[i].item())) for i in indices]


================================================================================
FILE: implementation/python/qubit/qft.py
================================================================================
"""Quantum Fourier Transform.

:func:`apply_qft` and :func:`apply_qft_inverse` implement the standard
Cooley-Tukey-style circuit (H + controlled-phase cascade + final
bit-reversal swaps), matching the ``/go`` and ``/c`` conventions:

* **LSB-first basis indexing.** Qubit 0 is the least significant bit
  of the basis index. Same as everywhere else in this package.
* **Natural binary order on output.** ``apply_qft`` includes the final
  SWAP cascade so the output amplitudes satisfy

      amp[y] = (1 / sqrt(N)) * exp(2 pi i * x * y / N)    for input |x>

  with ``N = 2**n`` and ``y`` indexed in the same flat-basis order as
  the input. Without the final swaps, the circuit naturally produces
  amplitudes at bit-reversed indices; the swaps undo that.

Loop ordering note: the outer loop processes qubits from
``start + n - 1`` (the high-index end of the sub-register) DOWN to
``start``. Within each iteration, the controlled-phase fan-in comes
from lower-indexed qubits with angles ``pi / 2**(i - j)``. This is
the convention that produces correct phases under LSB-first basis
indexing (a hand-trace against the analytic QFT|1> matches; an
alternative low-to-high outer loop produces the wrong phases on at
least n=2, x=1).

:func:`apply_qft_inverse` is a TRUE inverse, not three forward
applications (which would also work mathematically, since the QFT
satisfies F**4 == I, but at 3x the gate cost). The inverse reverses
the operation order, undoes the swaps first, then applies negated
controlled-phase angles in the reverse inner-loop order before each
H (H is self-inverse).
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

from ._assert import raise_value
from .gates_controlled import apply_controlled_phase, apply_swap
from .gates_single import apply_h

if TYPE_CHECKING:
    from .qreg import Qreg


def _resolve_n(q: Qreg, start: int, n: int | None) -> int:
    """Default ``n`` to "the rest of the register" and validate the range.

    Shared by :func:`apply_qft` and :func:`apply_qft_inverse` so both
    surface the same error messages for the same kind of misuse.
    """
    if n is None:
        n = q._n - start
    raise_value(start >= 0, "apply_qft: start=%d must be >= 0", start)
    raise_value(n >= 1, "apply_qft: n=%d must be >= 1", n)
    raise_value(
        start + n <= q._n,
        "apply_qft: start=%d + n=%d > q.n_qubits=%d",
        start, n, q._n,
    )
    return n


def apply_qft(q: Qreg, start: int = 0, n: int | None = None) -> None:
    """Apply the forward QFT to qubits ``[start, start + n)``.

    Defaults: ``start=0`` and ``n`` covering the rest of the register.

    Output amplitudes are in natural binary order (the standard
    bit-reversal swaps are included at the end). For input ``|x>``,
    the output satisfies ``amp[y] = exp(2 pi i x y / N) / sqrt(N)``
    where ``N = 2**n``.
    """
    n = _resolve_n(q, start, n)

    # Forward QFT: outer loop from the high-index end of the
    # sub-register down to the low. For each target qubit, apply H,
    # then the controlled-phase fan-in from every lower-indexed
    # qubit in the sub-register.
    for i in range(n - 1, -1, -1):
        target = start + i
        apply_h(q, target)
        for j in range(i - 1, -1, -1):
            control = start + j
            # The angle halves with each step further away from the
            # target: pi/2 for the nearest control, pi/4 next, etc.
            theta = math.pi / (1 << (i - j))
            apply_controlled_phase(q, control, target, theta)

    # Final bit-reversal swaps so output is in natural binary order.
    # Without these, amp[y] would land at bit_reverse(y).
    for i in range(n // 2):
        apply_swap(q, start + i, start + n - 1 - i)


def apply_qft_inverse(
    q: Qreg, start: int = 0, n: int | None = None
) -> None:
    """Apply the inverse QFT to qubits ``[start, start + n)``.

    True inverse (reversed gate order, negated phase angles), not
    three forward applications. ``apply_qft`` followed by
    ``apply_qft_inverse`` returns the original state on every basis
    input (within ``amp_tol_for(q.dtype)``).
    """
    n = _resolve_n(q, start, n)

    # Reverse the forward operation order: undo the swaps first
    # (SWAP is self-inverse, so we apply the same swap cascade), then
    # reverse the outer + inner loops with negated phase angles.
    for i in range(n // 2):
        apply_swap(q, start + i, start + n - 1 - i)

    for i in range(n):
        target = start + i
        # Inner loop is the inverse of the forward inner loop: range
        # from j=0 up to j=i-1, applying inverse controlled-phase
        # (negative angle) at each step. The order matters because
        # controlled-phase gates commute pairwise but the ENTIRE
        # inverse-QFT has to mirror the forward operation order.
        for j in range(i):
            control = start + j
            theta = -math.pi / (1 << (i - j))
            apply_controlled_phase(q, control, target, theta)
        apply_h(q, target)


================================================================================
FILE: implementation/python/qubit/grover.py
================================================================================
"""Grover's amplitude-amplification algorithm.

:func:`apply_grover` prepares the uniform superposition over the searched
qubits, then runs ``iterations`` rounds of (oracle + diffusion). The
oracle is a user-supplied callback that phase-flips marked basis states;
the diffusion is the standard ``H X MCZ X H`` sandwich (also called
"inversion about the mean").

Conventions and notes:

* **Search range.** The searched qubits are always ``[0, n_qubits)``,
  matching ``/c`` and ``/go``. A free-floating ``start`` parameter
  would generalise but isn't needed for v1 (Shor's period finder and
  every Grover application that ships here uses qubits-from-the-bottom).
* **Default iteration count.** When ``iterations=None`` we use
  ``floor(pi / 4 * sqrt(2 ** n_qubits))`` -- the textbook optimum for
  a single marked item. Callers with multiple marked items should pass
  ``iterations`` explicitly (e.g., one iteration for 4 marked out of
  16). The default is correct for the most common use case.
* **Global phase.** The diffusion as implemented (``H X MCZ X H``)
  differs from the textbook ``2|s><s| - I`` by an overall sign,
  introducing a ``(-1)**iterations`` global phase. Probabilities are
  unaffected -- tests assert on probabilities rather than raw
  amplitudes for this reason.
* **n_qubits == 1.** Supported but degenerate: the rotation angle
  ``arcsin(1/sqrt(2)) == pi/4`` means a single iteration over-rotates
  exactly to where you started, so ``P(marked) == 0.5`` for any
  iteration count. Documented and tested rather than rejected --
  rejecting would surprise callers who reduce ``n_qubits`` during
  experimentation.
"""

from __future__ import annotations

import math
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

from ._assert import raise_value
from .gates_multi import apply_multi_controlled_z
from .gates_single import apply_h, apply_x

if TYPE_CHECKING:
    from .qreg import Qreg


def apply_grover(
    q: Qreg,
    n_qubits: int,
    oracle: Callable[[Qreg, Any], None],
    user: Any = None,
    iterations: int | None = None,
) -> None:
    """Run Grover's algorithm on qubits ``[0, n_qubits)`` of ``q``.

    Applies ``H`` to each of the first ``n_qubits`` qubits and then
    iterates oracle + diffusion (``H^n X^n MCZ X^n H^n``).

    **Precondition:** qubits ``[0, n_qubits)`` must be in
    :math:`|0\\dots0\\rangle` on entry. The Hadamards then produce the
    uniform superposition the algorithm needs; if the register is in
    any other state, this routine still applies the same gate sequence
    but the meaning of the result is no longer "standard Grover." Call
    :meth:`Qreg.init_basis(0) <qubit.Qreg.init_basis>` first if in
    doubt. The remaining qubits (indices ``>= n_qubits``) are not
    touched and act as inert ancilla space.

    Args:
        q: register; ``q.n_qubits >= n_qubits``.
        n_qubits: number of qubits in the search space (the marked
            states live among the ``2 ** n_qubits`` basis states of
            qubits ``[0, n_qubits)``).
        oracle: callback ``(q, user) -> None`` that phase-flips marked
            basis states. Called once per iteration with the same
            ``user`` payload.
        user: opaque payload forwarded unchanged to every oracle call.
            Use it to thread the marked-state predicate (a single
            basis index, a list of indices, a hash function, etc.).
        iterations: how many Grover rounds to run. ``None`` (default)
            picks ``floor(pi/4 * sqrt(2 ** n_qubits))``, the optimum
            for a single marked item. Must be non-negative.

    Raises:
        ValueError: ``n_qubits`` out of ``[1, q.n_qubits]``, or
            ``iterations`` negative.
        TypeError: ``oracle`` is not callable.
    """
    raise_value(
        1 <= n_qubits <= q._n,
        "apply_grover: n_qubits=%d out of [1, %d]", n_qubits, q._n,
    )
    if not callable(oracle):
        raise TypeError(
            f"qubit: apply_grover: oracle must be callable, "
            f"got {type(oracle).__name__}"
        )

    if iterations is None:
        # Textbook optimal for a single marked item. For multiple
        # marked items the optimal count is smaller and callers
        # should pass `iterations` explicitly.
        n_search_states = 1 << n_qubits
        iterations = int(math.pi / 4.0 * math.sqrt(n_search_states))

    raise_value(
        iterations >= 0,
        "apply_grover: iterations=%d must be >= 0", iterations,
    )

    # Step 1: prepare the uniform superposition |s> = (1/sqrt(N)) sum_x |x>
    # over the searched qubits.
    searched = list(range(n_qubits))
    for i in searched:
        apply_h(q, i)

    # Step 2: Grover iterations. Each iteration:
    #   - oracle(q, user) phase-flips the marked basis states
    #   - diffusion (inversion about the mean) = H^n X^n MCZ_n X^n H^n
    for _ in range(iterations):
        oracle(q, user)
        for i in searched:
            apply_h(q, i)
        for i in searched:
            apply_x(q, i)
        apply_multi_controlled_z(q, searched)
        for i in searched:
            apply_x(q, i)
        for i in searched:
            apply_h(q, i)


================================================================================
FILE: implementation/python/qubit/shor.py
================================================================================
"""Shor's algorithm: modular exponentiation, period finding, factoring.

Three public entry points:

* :func:`apply_modular_exp` -- the reversible map
  ``(x, y) -> (x, (a**x * y) mod N)`` for ``y < N``, identity otherwise.
  Implemented as a single ``torch.gather`` with a precomputed integer
  permutation. The permutation is built on CPU because PyTorch's
  bitwise integer ops have inconsistent backend coverage (CUDA / ROCm /
  MPS); a contiguous device tensor is materialised once and gathered.
* :func:`apply_shor_period` -- the quantum period-finding subroutine.
  Initialises ``q`` to ``|0>_counting |1>_target``, runs the Shor
  circuit (``H^t -> modular_exp -> QFT^-1 -> measure``), and recovers
  a candidate period via continued-fraction expansion of ``c / 2^t``.
* :func:`shor_factor` -- end-to-end factoring with retry loop. Picks
  random ``a`` coprime to ``N``, invokes the quantum period finder,
  derives a factor via ``gcd(a**(r/2) +/- 1, N)``. Allocates its own
  ``Qreg``; size ``t + n`` with ``n = ceil(log2 N)``, ``t = 2n + 1``
  per the standard Shor recipe.

Scope note: v1 covers Shor-15 (13-qubit register: n=4 target bits +
t=2n+1=9 counting bits) and Shor-21 (16-qubit register: n=5+t=11)
comfortably -- both well under a second on CPU
including the CPU-side permutation build. 25-qubit modular
exponentiation is NOT a practical target for this implementation: the
permutation tensor would be 2^25 * 8 = 256 MiB on CPU before transfer
to device, and the gather output is another 2^25 * 16 = 512 MiB. The
arithmetic is correct at that scale, but call sites should consider
the memory cost.
"""

from __future__ import annotations

import dataclasses
import random
from typing import TYPE_CHECKING

import torch

from . import _memory
from ._assert import raise_value
from .gates_single import apply_h
from .qft import apply_qft_inverse
from .standart import continued_fraction, gcd_u64, mod_pow

if TYPE_CHECKING:
    from .qreg import Qreg


# ---------------------------------------------------------------------------
# Result dataclasses
# ---------------------------------------------------------------------------


@dataclasses.dataclass(frozen=True)
class ShorPeriodResult:
    """Result of one run of :func:`apply_shor_period`.

    Attributes:
        period: candidate period ``r`` recovered from the continued
            fraction expansion of ``measured / 2**t``. Zero if the
            recovery fell off the edge (degenerate measurement, unlikely
            for a well-sized counting register).
        measured: the raw measurement of the counting register (an
            integer in ``[0, 2**t)``). Exposed for debugging and for
            tests that want to inspect the underlying randomness rather
            than the recovered period.
    """

    period: int
    measured: int


@dataclasses.dataclass(frozen=True)
class ShorFactorResult:
    """Result of :func:`shor_factor`.

    Attributes:
        p, q: the two factors of ``N``. On success, ``p * q == N`` and
            both ``p, q > 1``. On failure (after exhausting
            ``max_attempts``), both fields are ``0``.
        attempts: number of period-finding rounds consumed. ``0`` for
            the trivial-even-``N`` short-circuit path; otherwise the
            attempt number on which a factor was found (or
            ``max_attempts`` on total failure).
    """

    p: int
    q: int
    attempts: int


# ---------------------------------------------------------------------------
# apply_modular_exp: the reversible (x, y) -> (x, a^x * y mod N) map.
# ---------------------------------------------------------------------------


def _build_modexp_perm(
    n_total: int,
    counting_start: int,
    t: int,
    target_start: int,
    n: int,
    a: int,
    N: int,
) -> torch.Tensor:
    """Build the gather permutation for one apply_modular_exp call.

    Returns a 1-D ``int64`` tensor of length ``2**n_total`` such that
    ``perm[i_new] == i_old`` for every input/output pair.

    The build is vectorised on CPU. Memory is one int64 tensor of size
    ``2**n_total``; for Shor-21 (n_total=16) that's 512 KiB, trivial.
    The ``pow(a, x, N)`` precompute is ``2**t`` Python int calls and
    runs in microseconds at our sizes.
    """
    size = 1 << n_total
    t_mask = (1 << t) - 1
    n_mask = (1 << n) - 1
    target_clear_mask = n_mask << target_start
    # Mask of bits to keep when reconstructing i_new: everything except
    # the target-register bits, within the n_total-bit window.
    keep_mask = (size - 1) & ~target_clear_mask

    # Precompute a^x mod N for x in [0, 2**t). Python's three-argument
    # pow is fast (C-level) and exact.
    a_pow_table = torch.tensor(
        [pow(a, x, N) for x in range(1 << t)],
        dtype=torch.int64,
    )

    idx = torch.arange(size, dtype=torch.int64)
    x = (idx >> counting_start) & t_mask
    y = (idx >> target_start) & n_mask

    # Use the precomputed table; advanced indexing into a 1-D tensor
    # is portable across backends.
    a_pow_x = a_pow_table[x]
    y_new_in_ring = (a_pow_x * y) % N
    # Pass-through for y >= N keeps the gate a bijection on the full
    # 2**n target space (required for reversibility / unitarity).
    y_new = torch.where(y < N, y_new_in_ring, y)

    i_new = (idx & keep_mask) | (y_new << target_start)

    # Scatter: perm[i_new[k]] = idx[k]. The map is a bijection because
    # gcd(a, N) == 1 makes y -> (a^x * y) mod N invertible on [0, N),
    # extended by identity on [N, 2^n). So i_new contains every value
    # in [0, size) exactly once and the fancy-indexing assignment is
    # well-defined (no duplicate destinations).
    perm = torch.empty(size, dtype=torch.int64)
    perm[i_new] = idx
    return perm


def apply_modular_exp(
    q: Qreg,
    counting_start: int,
    t: int,
    target_start: int,
    n: int,
    a: int,
    N: int,
) -> None:
    """Apply the reversible modular-exponentiation gate to ``q``.

    The counting register lives at qubits ``[counting_start, counting_start + t)``
    and stores ``x``. The target register lives at qubits
    ``[target_start, target_start + n)`` and stores ``y``. The action is

        (x, y) -> (x, (a**x * y) mod N)    for y < N
        (x, y) -> (x, y)                    for y >= N (pass-through)

    The two registers must not overlap, and ``gcd(a, N) == 1`` is
    required (otherwise the map is not a bijection and the operation
    isn't unitary).

    Implementation: builds the permutation tensor on CPU via
    :func:`_build_modexp_perm`, transfers it to ``q._device``, and
    applies a single ``torch.gather``. One state-vector allocation
    (the gather output) per call.
    """
    _validate_modexp_args(
        q, counting_start, t, target_start, n, a, N, caller="apply_modular_exp"
    )
    # Memory preflight for the permutation tensor + gather output, gated
    # by the Qreg's check_memory flag so power users can opt out the
    # same way they would on the constructor. Quiet no-op when the
    # device doesn't expose a free-memory query.
    if q._check_memory:
        _memory.preflight(q._n, q._dtype, q._device, op="modexp")

    perm = _build_modexp_perm(
        q._n, counting_start, t, target_start, n, a, N
    ).to(q._device)
    q._amp = q._amp.gather(0, perm)


def _validate_modexp_args(
    q: Qreg,
    counting_start: int,
    t: int,
    target_start: int,
    n: int,
    a: int,
    N: int,
    *,
    caller: str = "apply_modular_exp",
) -> None:
    """Validate the modular-exponentiation argument bundle.

    Shared by :func:`apply_modular_exp` and :func:`apply_shor_period`
    so that bad arguments are rejected BEFORE any state mutation. Each
    call site passes its own ``caller`` label so the raised
    ``ValueError`` messages stay precise.
    """
    raise_value(t >= 1, "%s: t=%d must be >= 1", caller, t)
    raise_value(n >= 1, "%s: n=%d must be >= 1", caller, n)
    raise_value(
        counting_start >= 0,
        "%s: counting_start=%d must be >= 0", caller, counting_start,
    )
    raise_value(
        target_start >= 0,
        "%s: target_start=%d must be >= 0", caller, target_start,
    )
    raise_value(
        counting_start + t <= q._n,
        "%s: counting range [%d, %d) out of [0, %d)",
        caller, counting_start, counting_start + t, q._n,
    )
    raise_value(
        target_start + n <= q._n,
        "%s: target range [%d, %d) out of [0, %d)",
        caller, target_start, target_start + n, q._n,
    )
    c_end = counting_start + t
    t_end = target_start + n
    raise_value(
        c_end <= target_start or t_end <= counting_start,
        "%s: counting [%d, %d) and target [%d, %d) overlap",
        caller, counting_start, c_end, target_start, t_end,
    )
    raise_value(N >= 2, "%s: N=%d must be >= 2", caller, N)
    raise_value(
        N <= 1 << n,
        "%s: N=%d exceeds target capacity 2^%d", caller, N, n,
    )
    raise_value(a >= 0, "%s: a=%d must be >= 0", caller, a)
    raise_value(
        gcd_u64(a, N) == 1,
        "%s: gcd(a=%d, N=%d) != 1 (not coprime)", caller, a, N,
    )


# ---------------------------------------------------------------------------
# apply_shor_period: one period-finding round.
# ---------------------------------------------------------------------------


def apply_shor_period(
    q: Qreg,
    counting_start: int,
    t: int,
    target_start: int,
    n: int,
    a: int,
    N: int,
) -> ShorPeriodResult:
    """Run one period-finding round of Shor's algorithm on ``q``.

    Re-initialises ``q`` (so the caller doesn't need to set it up
    first): the counting register is zeroed, the target register is set
    to ``|1>``, and any other qubits are zeroed. Then runs the standard
    Shor circuit:

    1. ``H`` on every counting qubit (uniform superposition over ``x``).
    2. :func:`apply_modular_exp` (correlates ``y`` with ``a**x mod N``).
    3. :func:`qubit.qft.apply_qft_inverse` on the counting register
       (concentrates probability mass at multiples of ``2**t / r``).
    4. Measure all qubits; extract the counting register's value.
    5. Recover candidate period via :func:`qubit.continued_fraction`
       applied to ``measured / 2**t`` with denominator bound ``N``.

    Caller validation: ``a`` and ``N`` must satisfy
    :func:`apply_modular_exp`'s preconditions (in particular
    ``gcd(a, N) == 1``). Counting and target ranges must not overlap.

    Returns a :class:`ShorPeriodResult`. The recovered ``period`` may be
    a divisor of the true order rather than the order itself; the
    caller (typically :func:`shor_factor`) is responsible for filtering
    odd / unlucky periods. A ``measured == 0`` readout (which carries
    no period information) returns ``period = 0`` so callers can treat
    "no useful period extracted" as a single sentinel.
    """
    # Validate arguments BEFORE init_basis() mutates the register, so
    # bad arguments don't leave the caller's register half-altered.
    # Same precondition bundle apply_modular_exp uses; passing
    # caller="apply_shor_period" so the error messages stay precise.
    _validate_modexp_args(
        q, counting_start, t, target_start, n, a, N,
        caller="apply_shor_period",
    )

    # Initialise state.
    q.init_basis(1 << target_start)

    # Step 1: uniform superposition over the counting register.
    for i in range(t):
        apply_h(q, counting_start + i)

    # Step 2: modular-exponentiation oracle.
    apply_modular_exp(q, counting_start, t, target_start, n, a, N)

    # Step 3: inverse QFT on the counting register.
    apply_qft_inverse(q, counting_start, t)

    # Step 4: measure. We measure ALL qubits and then bit-extract the
    # counting register's value, which is simpler than implementing a
    # sub-register measurement helper. The full measurement collapses
    # everything but Shor doesn't need the post-state -- the period
    # comes from the counting bits' value alone.
    full = q.measure_all()
    t_mask = (1 << t) - 1
    measured = (full >> counting_start) & t_mask

    # A c=0 readout carries no period information. continued_fraction
    # would return denominator 1 (Fraction(0, q).limit_denominator(...)
    # gives 0/1), which is a perfectly valid fraction but useless as a
    # period candidate. Report it as outright failure so callers don't
    # have to special-case period==1.
    if measured == 0:
        return ShorPeriodResult(period=0, measured=0)

    # Step 5: continued-fraction recovery of the candidate period.
    # Already shipping: continued_fraction(c, q, max_denom) returns the
    # denominator of the best rational approximation of c/q with
    # denominator <= max_denom. For Shor that bound is N (the modulus
    # being factored).
    period = continued_fraction(measured, 1 << t, N)
    return ShorPeriodResult(period=period, measured=measured)


# ---------------------------------------------------------------------------
# shor_factor: end-to-end factoring with retry loop.
# ---------------------------------------------------------------------------


def shor_factor(
    N: int,
    max_attempts: int = 20,
    seed: int | None = None,
) -> ShorFactorResult:
    """Factor ``N`` end-to-end via Shor's algorithm.

    Args:
        N: the number to factor. Must be >= 2.
        max_attempts: maximum number of period-finding rounds before
            declaring failure. Default 20 is comfortable for Shor-15;
            ~10 attempts is the realistic median success threshold.
        seed: optional integer for reproducibility. Threads through to
            both the Python ``random.Random`` driving the choice of
            ``a`` AND each constructed Qreg's measurement RNG, so
            seeded runs are bit-for-bit reproducible.

    Returns a :class:`ShorFactorResult`. On the trivial-even-``N``
    short-circuit ``attempts == 0``; otherwise ``attempts`` is the
    attempt number on which a factor was found, or ``max_attempts`` on
    failure.

    Scope caveat: this function does NOT detect prime powers
    (``N = p**k`` for prime ``p``, ``k >= 2``). For those inputs, the
    period-finding step has no useful order to recover, every attempt
    fails the ``a**(r/2) != N - 1`` check or returns ``period == 0``,
    and the function exhausts ``max_attempts`` then reports failure.
    The classical pre-check (try ``round(N**(1/k)) ** k == N`` for
    every ``k`` in ``[2, log2 N]``) belongs in the caller, matching
    ``/c`` and ``/go``'s scope. Even ``N`` is handled here (short-
    circuit).

    Algorithm:

    1. Validate ``N >= 2`` and ``max_attempts >= 1``.
    2. If ``N`` is even, return ``(2, N // 2)`` without invoking the
       quantum subroutine.
    3. Otherwise pick a random ``a`` in ``[2, N - 1]``. If
       ``gcd(a, N) > 1`` we already have a non-trivial factor.
    4. Run :func:`apply_shor_period` with that ``a`` on a freshly
       allocated ``Qreg(t + n)`` where ``n = ceil(log2 N)`` and
       ``t = 2n + 1`` (standard Shor recipe).
    5. Check the recovered period ``r``: it must be non-zero and even,
       and ``a**(r/2) mod N`` must not equal ``N - 1`` (otherwise the
       gcd step yields only trivial factors).
    6. Compute ``gcd(a**(r/2) +/- 1, N)``; if either gives a
       non-trivial factor, return it.
    7. On failure, loop back to step 3 with a fresh ``a`` and a fresh
       Qreg.
    """
    raise_value(N >= 2, "shor_factor: N=%d must be >= 2", N)
    raise_value(
        max_attempts >= 1,
        "shor_factor: max_attempts=%d must be >= 1", max_attempts,
    )

    # Even-N short-circuit. Returns sorted (small, large) for
    # determinism.
    if N % 2 == 0:
        return ShorFactorResult(p=2, q=N // 2, attempts=0)

    # Register sizing: n is the smallest number of bits that can
    # represent every value in [0, N). For N=15 that's 4 (since
    # 14 = 0b1110). For N=21 that's 5 (since 20 = 0b10100). The
    # counting register gets t = 2n + 1 bits, the standard Shor recipe
    # that guarantees >= 1/2 success probability per attempt.
    n = (N - 1).bit_length()
    t_bits = 2 * n + 1
    n_total = t_bits + n

    # The python `random.Random(seed)` drives BOTH the choice of `a`
    # AND the seed for each constructed Qreg's measurement RNG. Seeded
    # runs are bit-for-bit reproducible.
    rng = random.Random(seed)

    # Lazy import: shor.py is at the same level as qreg.py and qreg.py
    # has method wrappers that import shor lazily. Doing it the other
    # direction with TYPE_CHECKING + lazy here keeps both directions
    # cycle-free at module load.
    from .qreg import Qreg

    for attempt in range(1, max_attempts + 1):
        a = rng.randint(2, N - 1)
        # Derive a Qreg seed from the same RNG so the full call is
        # reproducible given just `seed`. Use a wide enough range to
        # exercise the full 64-bit measurement RNG space.
        qreg_seed = rng.randint(0, (1 << 63) - 1)

        # Lucky-shortcut: if a is not coprime to N, gcd itself gives a
        # non-trivial factor (because gcd(a, N) divides N and 1 < gcd
        # < N when gcd != 1).
        g = gcd_u64(a, N)
        if g != 1:
            return ShorFactorResult(p=g, q=N // g, attempts=attempt)

        # Allocate a fresh register; allow seeding for reproducibility.
        # We leave the default check_memory=True. The supported v1 cases
        # are small (Shor-15 = 13 qubits, Shor-21 = 16) so the preflight
        # is essentially free, and apply_modular_exp also preflights its
        # own permutation+gather peak before allocating.
        q = Qreg(n_total, seed=qreg_seed)
        res = apply_shor_period(q, n, t_bits, 0, n, a, N)
        r = res.period

        if r == 0 or r % 2 != 0:
            continue

        # Verify r is actually a period of a mod N. Continued-fraction
        # recovery can return a divisor of the true order; the gcd
        # extraction below would just produce trivial factors in that
        # case and we'd retry, but it's cheaper to filter upfront and
        # matches the verification promise documented on apply_shor_period.
        if mod_pow(a, r, N) != 1:
            continue

        half = mod_pow(a, r // 2, N)
        if half == N - 1:
            # a^(r/2) ≡ -1 (mod N) -- the gcd step would only yield
            # trivial factors (1 or N). Retry with a different `a`.
            continue

        # gcd(a^(r/2) + 1, N) or gcd(a^(r/2) - 1, N) gives a factor.
        candidate = gcd_u64(half + 1, N)
        if 1 < candidate < N:
            return ShorFactorResult(
                p=candidate, q=N // candidate, attempts=attempt
            )
        if half > 0:
            candidate = gcd_u64(half - 1, N)
            if 1 < candidate < N:
                return ShorFactorResult(
                    p=candidate, q=N // candidate, attempts=attempt
                )

    return ShorFactorResult(p=0, q=0, attempts=max_attempts)


================================================================================
FILE: implementation/python/qubit/cli.py
================================================================================
"""Command-line demo for the qubit library.

Installed as ``qubit-demo`` via ``[project.scripts]`` in
:file:`pyproject.toml`. Mirrors the demos in ``/c``'s ``cmd/qubit`` and
``/go``'s ``cmd/qubit/main.go``: a quick smoke check that the library
runs end-to-end on whatever device the user has, for each of the four
canonical algorithm-level demos (Bell state, QFT, Grover, Shor).

Run from a checkout of the repo::

    cd implementation/python
    uv sync --group dev    # one-time install
    uv run qubit-demo --algo bell
    uv run qubit-demo --algo qft
    uv run qubit-demo --algo grover
    uv run qubit-demo --algo shor

All demos are deterministic (seeded RNG where used) so successive runs
produce identical output, and they exit with status 0 on success. Any
:class:`qubit:`-prefixed exception or unexpected error is printed to
stderr and the process exits non-zero.
"""

from __future__ import annotations

import argparse
import math
import sys
import warnings
from collections.abc import Callable
from typing import Any

# PyTorch's optional NumPy bridge warns on import when NumPy is absent.
# We don't depend on NumPy, so the warning is pure noise in demo output.
# Filter it before importing the modules that pull torch in.
warnings.filterwarnings(
    "ignore",
    message="Failed to initialize NumPy",
    category=UserWarning,
)

from .grover import apply_grover  # noqa: E402  -- after warning filter
from .qreg import Qreg  # noqa: E402
from .shor import shor_factor  # noqa: E402


def _demo_bell() -> None:
    """Bell state ``(|00> + |11>) / sqrt(2)`` from H + CNOT."""
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    q.apply_h(0)
    q.apply_cnot(0, 1)
    print(
        f"Bell |Phi+>: P(00) = {q.prob_of(0):.4f}, "
        f"P(11) = {q.prob_of(3):.4f}"
    )


def _demo_qft() -> None:
    """``QFT|0>`` on a 4-qubit register -- uniform distribution."""
    n = 4
    q = Qreg(n, device="cpu")
    q.init_basis(0)
    q.apply_qft()
    expected = 1.0 / (1 << n)
    print(
        f"QFT|0> on {n} qubits: "
        f"P(0) = {q.prob_of(0):.4f} (uniform = {expected:.4f})"
    )


def _demo_grover() -> None:
    """Grover finds ``|1111>`` on a 4-qubit register.

    Uses :func:`apply_multi_controlled_z` as the oracle (which marks
    ``|1...1>`` over the searched qubits). The default iteration count
    is the textbook optimum for one marked item among 16 -- 3 rounds,
    giving an analytic recovery probability of ~0.961.
    """
    n = 4
    q = Qreg(n, device="cpu")
    q.init_basis(0)

    def oracle(q_inner: Qreg, _user: Any) -> None:
        # Mark |1...1> by phase-flipping the all-ones amplitude.
        q_inner.apply_multi_controlled_z(list(range(n)))

    iters = int(math.pi / 4.0 * math.sqrt(1 << n))
    apply_grover(q, n, oracle, iterations=iters)
    print(
        f"Grover marked |1{'1' * (n - 1)}>: "
        f"P({(1 << n) - 1}) = {q.prob_of((1 << n) - 1):.4f}"
    )


def _demo_shor() -> None:
    """Factor ``N = 15`` end-to-end via :func:`shor_factor`.

    Seeded for determinism so this demo is reproducible across runs.
    """
    res = shor_factor(15, seed=42)
    print(
        f"Shor(15, seed=42): p={res.p}, q={res.q}, "
        f"attempts={res.attempts}"
    )


_DEMOS: dict[str, Callable[[], None]] = {
    "bell": _demo_bell,
    "qft": _demo_qft,
    "grover": _demo_grover,
    "shor": _demo_shor,
}


def main(argv: list[str] | None = None) -> int:
    """Run the ``qubit-demo`` CLI. Returns 0 on success, 1 on error.

    Argparse handles ``--help`` and unknown options; both exit through
    ``SystemExit`` rather than our exception path. Demo errors go to
    stderr with a non-zero return.
    """
    parser = argparse.ArgumentParser(
        prog="qubit-demo",
        description=(
            "Run a quick algorithm-level demo of the qubit library. "
            "Mirrors the bell/qft/grover/shor demos in /c and /go."
        ),
    )
    parser.add_argument(
        "--algo",
        choices=sorted(_DEMOS),
        default="bell",
        help="Which algorithm to demo (default: bell).",
    )
    args = parser.parse_args(argv)

    try:
        _DEMOS[args.algo]()
    except Exception as exc:  # noqa: BLE001 -- intentional catch-all at CLI top.
        print(f"qubit-demo: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())


================================================================================
FILE: implementation/python/tests/__init__.py
================================================================================


================================================================================
FILE: implementation/python/tests/conftest.py
================================================================================
"""Pytest fixtures shared across the test suite.

The ``device`` fixture parametrises tests across whatever devices the
running machine has available. On CPU-only CI it produces just
``"cpu"``; on Apple Silicon, ``["cpu", "mps"]``; on a CUDA box,
``["cpu", "cuda"]``. Tests that don't care about device just don't
request the fixture.
"""

from __future__ import annotations

import pytest
import torch


def _available_devices() -> list[str]:
    """List of device strings PyTorch can use on this host.

    CPU is always present. CUDA and MPS are added if their respective
    runtime detections succeed. Evaluated once at test-collection time.
    """
    devs = ["cpu"]
    if torch.cuda.is_available():
        devs.append("cuda")
    if torch.backends.mps.is_available():
        devs.append("mps")
    return devs


_DEVICES = _available_devices()


@pytest.fixture(params=_DEVICES, ids=lambda d: d)
def device(request: pytest.FixtureRequest) -> str:
    """A device string for every device PyTorch can reach on this host."""
    param: str = request.param
    return param


================================================================================
FILE: implementation/python/tests/test_import.py
================================================================================
"""Phase-0 smoke test: the package imports and re-exports its public surface."""

from __future__ import annotations


def test_package_imports() -> None:
    import qubit

    # Public surface promised by __init__.py.
    assert hasattr(qubit, "Qreg")
    assert hasattr(qubit, "qubit_axis")
    assert hasattr(qubit, "amp_tol_for")
    assert hasattr(qubit, "prob_tol_for")
    for name in (
        "AMP_TOL_C64",
        "AMP_TOL_C128",
        "PROB_TOL_C64",
        "PROB_TOL_C128",
    ):
        assert hasattr(qubit, name), f"missing public constant {name}"


def test_torch_available() -> None:
    # Sanity-check the only third-party dependency.
    import torch

    # Any modern PyTorch is fine for Phase 0+1; gate code (Phase 4+) will
    # exercise the complex-tensor surface that requires >= 2.4 in practice.
    assert torch.__version__.split(".")[0] in {"2", "3"}


================================================================================
FILE: implementation/python/tests/test_assert.py
================================================================================
"""Tests for the package-private exception helpers."""

from __future__ import annotations

import pytest

from qubit._assert import raise_type, raise_value


def test_raise_value_passes_when_true() -> None:
    raise_value(True, "should not fire")


def test_raise_value_raises_with_prefix() -> None:
    with pytest.raises(ValueError) as exc_info:
        raise_value(False, "bad: %d", 42)
    msg = str(exc_info.value)
    assert msg.startswith("qubit: ")
    assert "bad: 42" in msg


def test_raise_value_formats_multiple_args() -> None:
    with pytest.raises(ValueError, match=r"qubit: range \[5, 10\)"):
        raise_value(False, "range [%d, %d)", 5, 10)


def test_raise_type_passes_when_true() -> None:
    raise_type(True, "should not fire")


def test_raise_type_raises_with_prefix() -> None:
    with pytest.raises(TypeError) as exc_info:
        raise_type(False, "expected %s, got %s", "Tensor", "list")
    msg = str(exc_info.value)
    assert msg.startswith("qubit: ")
    assert "expected Tensor, got list" in msg


def test_raise_value_and_raise_type_use_different_classes() -> None:
    # A test that catches ValueError must NOT catch the TypeError version.
    with pytest.raises(TypeError):
        raise_type(False, "should be TypeError")
    with pytest.raises(ValueError):
        raise_value(False, "should be ValueError")


================================================================================
FILE: implementation/python/tests/test_axis.py
================================================================================
"""Tests for the qubit-to-tensor-axis convention.

The whole package's correctness depends on the LSB-first mapping
``axis = n - 1 - target``. These tests lock the convention so any future
refactor that flips qubit ordering fails loudly here before the gate
tests start producing mysteriously bit-reversed outputs.
"""

from __future__ import annotations

import pytest

from qubit import qubit_axis


def test_qubit_0_is_last_axis() -> None:
    # Qubit 0 is LSB -> rightmost tensor axis when state is reshaped to (2,)*n.
    assert qubit_axis(0, 1) == 0
    assert qubit_axis(0, 2) == 1
    assert qubit_axis(0, 4) == 3
    assert qubit_axis(0, 10) == 9


def test_qubit_n_minus_1_is_first_axis() -> None:
    # The MSB lives at axis 0 (slowest-varying axis).
    assert qubit_axis(0, 1) == 0  # n=1: only one axis, sits at index 0
    assert qubit_axis(1, 2) == 0
    assert qubit_axis(3, 4) == 0
    assert qubit_axis(9, 10) == 0


def test_full_axis_sweep_is_a_permutation() -> None:
    # For every n, the set of qubit_axis values must be a permutation
    # of [0, n) -- no qubit collides with another, no axis goes unused.
    for n in (1, 2, 3, 4, 8, 16):
        axes = [qubit_axis(q, n) for q in range(n)]
        assert sorted(axes) == list(range(n))


def test_axis_mapping_is_self_inverse() -> None:
    # qubit_axis is its own inverse: axis(axis(q, n), n) == q.
    for n in (1, 4, 8):
        for q in range(n):
            assert qubit_axis(qubit_axis(q, n), n) == q


@pytest.mark.parametrize(
    "target,n",
    [(-1, 1), (1, 1), (4, 4), (100, 4), (-5, 3)],
)
def test_out_of_range_raises(target: int, n: int) -> None:
    with pytest.raises(ValueError, match=r"^qubit: qubit_axis: "):
        qubit_axis(target, n)


================================================================================
FILE: implementation/python/tests/test_device.py
================================================================================
"""Tests for the device-selection and dtype-policy module.

These cover the MPS+complex128 ban that is critical to avoid silent
downgrades at construction time. Tests that need a specific device skip
gracefully when the device isn't available on the host running the suite.
"""

from __future__ import annotations

import pytest
import torch

from qubit._device import (
    coerce_device,
    default_device,
    default_dtype,
    validate_dtype_device,
)


def test_default_device_is_torch_device() -> None:
    d = default_device()
    assert isinstance(d, torch.device)
    # Must be one of the three kinds we support.
    assert d.type in {"cpu", "cuda", "mps"}


def test_default_dtype_on_cpu() -> None:
    assert default_dtype(torch.device("cpu")) == torch.complex128


def test_default_dtype_on_mps() -> None:
    # The policy holds regardless of whether MPS is actually available
    # on the host: the function maps device type -> dtype, no probe.
    assert default_dtype(torch.device("mps")) == torch.complex64


def test_default_dtype_on_cuda() -> None:
    # CUDA defaults to complex128 like CPU; same policy applies for ROCm
    # (also reported as device.type == "cuda").
    assert default_dtype(torch.device("cuda")) == torch.complex128


def test_validate_cpu_complex128_ok() -> None:
    # Should not raise.
    validate_dtype_device(torch.device("cpu"), torch.complex128)


def test_validate_cpu_complex64_ok() -> None:
    validate_dtype_device(torch.device("cpu"), torch.complex64)


def test_validate_mps_complex64_ok() -> None:
    validate_dtype_device(torch.device("mps"), torch.complex64)


def test_validate_mps_complex128_raises() -> None:
    with pytest.raises(ValueError, match=r"^qubit: MPS backend does not support complex128"):
        validate_dtype_device(torch.device("mps"), torch.complex128)


def test_validate_rejects_non_complex_dtype() -> None:
    with pytest.raises(ValueError, match=r"^qubit: validate_dtype_device"):
        validate_dtype_device(torch.device("cpu"), torch.float64)


def test_coerce_none_returns_default() -> None:
    out = coerce_device(None)
    assert out == default_device()


def test_coerce_string() -> None:
    assert coerce_device("cpu") == torch.device("cpu")


def test_coerce_device_passthrough() -> None:
    d = torch.device("cpu")
    assert coerce_device(d) is d


================================================================================
FILE: implementation/python/tests/test_memory.py
================================================================================
"""Tests for the memory-preflight helpers."""

from __future__ import annotations

import pytest
import torch

from qubit._memory import (
    dtype_bytes,
    estimate_peak_bytes,
    estimate_state_bytes,
    free_bytes,
    preflight,
)


def test_dtype_bytes_complex64() -> None:
    assert dtype_bytes(torch.complex64) == 8


def test_dtype_bytes_complex128() -> None:
    assert dtype_bytes(torch.complex128) == 16


def test_dtype_bytes_unsupported() -> None:
    with pytest.raises(ValueError, match=r"^qubit: dtype_bytes: "):
        dtype_bytes(torch.float64)


def test_estimate_state_bytes_minimum() -> None:
    # 1 qubit -> 2 amplitudes -> 32 bytes at complex128, 16 at complex64.
    assert estimate_state_bytes(1, torch.complex128) == 32
    assert estimate_state_bytes(1, torch.complex64) == 16


def test_estimate_state_bytes_at_25_qubits() -> None:
    # The /go thesis target; 512 MiB at complex128.
    assert estimate_state_bytes(25, torch.complex128) == (1 << 25) * 16
    assert estimate_state_bytes(25, torch.complex64) == (1 << 25) * 8


def test_estimate_state_bytes_rejects_zero() -> None:
    with pytest.raises(ValueError, match=r"^qubit: estimate_state_bytes: "):
        estimate_state_bytes(0, torch.complex128)


def test_estimate_state_bytes_rejects_negative() -> None:
    with pytest.raises(ValueError):
        estimate_state_bytes(-3, torch.complex128)


def test_peak_state_equals_state_bytes() -> None:
    # The "state" op is the baseline; no extra allocation.
    n = 10
    dt = torch.complex128
    assert estimate_peak_bytes(n, dt, op="state") == estimate_state_bytes(n, dt)


def test_peak_qft_is_four_state_bytes() -> None:
    # QFT is NOT actually in-place: it decomposes into Hadamards and
    # controlled-phase gates, each of which allocates a fresh state
    # vector via PyTorch tensor ops. We treat 4*state as a defensive
    # upper bound on a single QFT gate's transient peak.
    n = 10
    dt = torch.complex128
    assert estimate_peak_bytes(n, dt, op="qft") == 4 * estimate_state_bytes(n, dt)


def test_peak_single_gate_is_two_state_bytes() -> None:
    # apply_u runs tensordot + movedim, which holds the source state
    # vector and the freshly-allocated output simultaneously.
    n = 10
    dt = torch.complex128
    assert estimate_peak_bytes(n, dt, op="single_gate") == 2 * estimate_state_bytes(n, dt)


def test_peak_controlled_gate_is_four_state_bytes() -> None:
    # apply_cu does permute + reshape + matmul + reshape + inverse-permute;
    # the transient peak can hold up to four state-sized tensors.
    n = 10
    dt = torch.complex128
    assert estimate_peak_bytes(n, dt, op="controlled_gate") == 4 * estimate_state_bytes(n, dt)


def test_peak_modexp_counts_extra_tensors() -> None:
    # ModularExp peak = 2*state (gather output + original) + 2*int64
    # permutation (CPU side + device side transiently coexist).
    n = 10
    dt = torch.complex128
    state = estimate_state_bytes(n, dt)
    perm = (1 << n) * 8
    assert estimate_peak_bytes(n, dt, op="modexp") == 2 * state + 2 * perm


def test_peak_unknown_op_raises() -> None:
    # Deliberate: passing a string outside the _Op Literal union to exercise
    # the runtime branch. The `type: ignore` on the call itself suppresses
    # mypy's complaint about the bad Literal argument.
    with pytest.raises(ValueError, match=r"^qubit: estimate_peak_bytes: unknown op"):
        estimate_peak_bytes(4, torch.complex128, op="bogus")  # type: ignore[arg-type]


def test_free_bytes_cpu_returns_none() -> None:
    # CPU has no portable free-memory query in stable PyTorch, so the
    # function reports None to let preflight skip the precise check.
    assert free_bytes(torch.device("cpu")) is None


def test_free_bytes_mps_returns_none() -> None:
    assert free_bytes(torch.device("mps")) is None


def test_preflight_accepts_small_request() -> None:
    # Small requests under the sanity ceiling pass without raising on
    # any device (free_bytes is None on CPU/MPS so the precise check
    # is skipped; the CUDA branch has plenty of headroom for 8 qubits).
    preflight(8, torch.complex128, torch.device("cpu"), op="state")


def test_preflight_rejects_obviously_impossible_request() -> None:
    # 50 qubits at complex128 = 18 PiB. Far above the 1-TiB sanity
    # ceiling regardless of device. Must raise MemoryError before any
    # allocation is attempted.
    with pytest.raises(MemoryError, match=r"^qubit: preflight: requested n_qubits=50"):
        preflight(50, torch.complex128, torch.device("cpu"), op="state")


def test_preflight_rejects_modexp_oversize() -> None:
    # ModularExp at 40 qubits would need ~32 TiB peak; reject.
    with pytest.raises(MemoryError):
        preflight(40, torch.complex128, torch.device("cpu"), op="modexp")


================================================================================
FILE: implementation/python/tests/test_qreg.py
================================================================================
"""Tests for the Qreg class -- construction, properties, and accessors.

These exercise every public method that ships in Phase 0+1. Gate
methods, measurement, QFT, Grover, and Shor land in later phases and
have their own test files.
"""

from __future__ import annotations

import pytest
import torch

from qubit import (
    AMP_TOL_C64,
    AMP_TOL_C128,
    PROB_TOL_C64,
    PROB_TOL_C128,
    Qreg,
    amp_tol_for,
    prob_tol_for,
)

# ---------------------------------------------------------------------------
# Construction: valid and invalid inputs
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("n", [1, 2, 3, 4, 8])
def test_construct_valid_n(n: int) -> None:
    q = Qreg(n)
    assert q.n_qubits == n
    assert len(q._amp) == (1 << n)
    # A freshly-constructed Qreg is in the computational basis |0...0>:
    # amp[0] = 1, all other amplitudes 0, norm 1. This makes the
    # register a valid quantum state from the moment NewQreg returns,
    # so callers can apply gates immediately without an explicit
    # init_basis(0) (which still exists for explicit resets).
    assert q._amp[0].item() == complex(1, 0)
    assert q.norm() == pytest.approx(1.0)
    for i in range(1, 1 << n):
        assert q._amp[i].item() == complex(0, 0)


def test_construct_rejects_n_zero() -> None:
    with pytest.raises(ValueError, match=r"^qubit: Qreg: n_qubits=0"):
        Qreg(0)


def test_construct_rejects_n_negative() -> None:
    with pytest.raises(ValueError, match=r"^qubit: Qreg: n_qubits=-3"):
        Qreg(-3)


def test_construct_rejects_non_int_n() -> None:
    with pytest.raises(TypeError, match=r"^qubit: Qreg: n_qubits must be int"):
        Qreg(3.5)  # type: ignore[arg-type]


def test_construct_rejects_bool_n() -> None:
    # `bool` is a subclass of `int` in Python (so mypy sees Qreg(True) as
    # type-compatible), but we reject it explicitly at runtime so
    # Qreg(True) doesn't silently make a 1-qubit register.
    with pytest.raises(TypeError, match=r"^qubit: Qreg: n_qubits must be int"):
        Qreg(True)


def test_construct_rejects_oversize_obvious() -> None:
    # Past the 1 TiB sanity ceiling -> MemoryError before allocation.
    with pytest.raises(MemoryError, match=r"^qubit: preflight: "):
        Qreg(50)


def test_construct_check_memory_false_skips_preflight() -> None:
    # The escape hatch should let through requests that the preflight
    # would reject. We still don't actually allocate (would OOM) -- just
    # confirm the MemoryError doesn't come from the preflight layer by
    # catching whatever runtime error PyTorch raises instead.
    with pytest.raises(Exception) as exc_info:
        Qreg(50, check_memory=False)
    # The error MUST NOT be the preflight's MemoryError with our prefix.
    msg = str(exc_info.value)
    assert not msg.startswith("qubit: preflight:"), (
        f"check_memory=False should skip preflight; got our message: {msg}"
    )


# ---------------------------------------------------------------------------
# Device / dtype resolution
# ---------------------------------------------------------------------------


def test_explicit_cpu_device() -> None:
    q = Qreg(3, device="cpu")
    assert q.device.type == "cpu"
    assert q.dtype == torch.complex128


def test_explicit_cpu_complex64() -> None:
    q = Qreg(3, device="cpu", dtype=torch.complex64)
    assert q.dtype == torch.complex64


def test_default_device_matches_helper() -> None:
    from qubit._device import default_device

    q = Qreg(3)
    assert q.device == default_device()


def test_mps_complex128_rejected() -> None:
    # The MPS-banned-dtype check runs regardless of whether MPS is
    # available -- validate_dtype_device is purely a policy function,
    # so this test is portable to CPU-only CI.
    with pytest.raises(ValueError, match=r"MPS backend does not support complex128"):
        Qreg(3, device="mps", dtype=torch.complex128)


@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable")
def test_mps_default_dtype_is_complex64() -> None:
    q = Qreg(3, device="mps")
    assert q.dtype == torch.complex64
    assert q.device.type == "mps"


@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable")
def test_mps_explicit_complex64_ok() -> None:
    q = Qreg(3, device="mps", dtype=torch.complex64)
    assert q.dtype == torch.complex64


# ---------------------------------------------------------------------------
# Seed and RNG
# ---------------------------------------------------------------------------


def test_construct_with_seed() -> None:
    q1 = Qreg(2, seed=42)
    q2 = Qreg(2, seed=42)
    # Both generators were seeded with the same int; their next draws
    # must match exactly. This validates that seed= goes through to the
    # RNG and that the RNG lives on CPU (so seeded reproducibility holds
    # regardless of where _amp ends up).
    draw1 = torch.rand((), generator=q1._gen).item()
    draw2 = torch.rand((), generator=q2._gen).item()
    assert draw1 == draw2


def test_construct_without_seed_diversifies() -> None:
    # Two unseeded Qregs constructed back-to-back should produce
    # independent RNG sequences. Not a strict guarantee (extremely
    # unlikely collision is possible) but the practical contract.
    q1 = Qreg(2)
    q2 = Qreg(2)
    draw1 = torch.rand((), generator=q1._gen).item()
    draw2 = torch.rand((), generator=q2._gen).item()
    assert draw1 != draw2


def test_rng_is_cpu_generator() -> None:
    # MPS-generator quirks motivate keeping the RNG on CPU even when
    # _amp lives on GPU. The generator's device should always be CPU.
    q = Qreg(2)
    assert q._gen.device.type == "cpu"


# ---------------------------------------------------------------------------
# init_basis
# ---------------------------------------------------------------------------


def test_init_basis_sets_target_amp_to_one() -> None:
    q = Qreg(3)
    q.init_basis(5)
    # |101> -> amp[5] = 1 + 0j, all else zero.
    for i in range(8):
        want = 1 + 0j if i == 5 else 0 + 0j
        assert q.amplitude(i) == want


def test_init_basis_zero_state() -> None:
    q = Qreg(3)
    q.init_basis(0)
    assert q.amplitude(0) == 1 + 0j
    for i in range(1, 8):
        assert q.amplitude(i) == 0 + 0j


def test_init_basis_zeros_previous_state() -> None:
    # init_basis must clear whatever was there before.
    q = Qreg(3)
    q.init_basis(7)
    q.init_basis(2)
    assert q.amplitude(2) == 1 + 0j
    assert q.amplitude(7) == 0 + 0j


def test_init_basis_rejects_out_of_range() -> None:
    q = Qreg(3)
    with pytest.raises(ValueError, match=r"^qubit: init_basis: basis=8"):
        q.init_basis(8)


def test_init_basis_rejects_negative() -> None:
    q = Qreg(3)
    with pytest.raises(ValueError, match=r"^qubit: init_basis: basis=-1"):
        q.init_basis(-1)


# ---------------------------------------------------------------------------
# amplitude
# ---------------------------------------------------------------------------


def test_amplitude_reads_back_one() -> None:
    q = Qreg(2)
    q.init_basis(2)
    assert q.amplitude(2) == 1 + 0j


def test_amplitude_reads_back_zero() -> None:
    q = Qreg(2)
    q.init_basis(2)
    assert q.amplitude(3) == 0 + 0j


def test_amplitude_returns_complex_type() -> None:
    q = Qreg(2)
    q.init_basis(0)
    assert isinstance(q.amplitude(0), complex)


def test_amplitude_rejects_out_of_range() -> None:
    q = Qreg(2)
    with pytest.raises(ValueError, match=r"^qubit: amplitude: i=4"):
        q.amplitude(4)


def test_amplitude_rejects_negative() -> None:
    q = Qreg(2)
    with pytest.raises(ValueError, match=r"^qubit: amplitude: i=-1"):
        q.amplitude(-1)


# ---------------------------------------------------------------------------
# amplitudes_copy
# ---------------------------------------------------------------------------


def test_amplitudes_copy_is_cpu_tensor() -> None:
    # Even with no GPU available, the contract is "always CPU"; verify
    # that explicitly so the same assertion holds when run on MPS.
    q = Qreg(3)
    copy = q.amplitudes_copy()
    assert copy.device.type == "cpu"


def test_amplitudes_copy_independent_of_qreg() -> None:
    q = Qreg(2)
    q.init_basis(0)
    copy = q.amplitudes_copy()
    # Mutate the copy; the live amp must be unaffected.
    copy[0] = 99 + 0j
    assert q.amplitude(0) == 1 + 0j


def test_amplitudes_copy_length_and_dtype() -> None:
    q = Qreg(3)
    copy = q.amplitudes_copy()
    assert copy.shape == (8,)
    assert copy.dtype == q.dtype


def test_amplitudes_copy_reflects_state_at_call_time() -> None:
    q = Qreg(2)
    q.init_basis(1)
    snap1 = q.amplitudes_copy()
    q.init_basis(2)
    snap2 = q.amplitudes_copy()
    # snap1 captured |01>; snap2 captured |10>. They must differ.
    assert not torch.equal(snap1, snap2)
    assert snap1[1].item() == 1 + 0j
    assert snap2[2].item() == 1 + 0j


# ---------------------------------------------------------------------------
# prob_of and norm
# ---------------------------------------------------------------------------


def test_prob_of_basis_state_one() -> None:
    q = Qreg(2)
    q.init_basis(3)
    assert q.prob_of(3) == pytest.approx(1.0, abs=PROB_TOL_C128)


def test_prob_of_basis_state_zero() -> None:
    q = Qreg(2)
    q.init_basis(3)
    assert q.prob_of(0) == pytest.approx(0.0, abs=PROB_TOL_C128)


def test_prob_of_rejects_out_of_range() -> None:
    q = Qreg(2)
    with pytest.raises(ValueError, match=r"^qubit: prob_of: basis=4"):
        q.prob_of(4)


def test_norm_basis_state_is_one() -> None:
    q = Qreg(4)
    q.init_basis(7)
    assert q.norm() == pytest.approx(1.0, abs=PROB_TOL_C128)


def test_norm_zero_state_is_zero() -> None:
    # A *fresh* Qreg is now in |0...0> with norm 1. Explicitly zero
    # the amplitude vector to exercise the zero-norm path; this keeps
    # the test exercising norm()'s behavior on the pathological-but-
    # well-defined zero state.
    q = Qreg(3)
    q._amp.zero_()
    assert q.norm() == pytest.approx(0.0, abs=PROB_TOL_C128)


# ---------------------------------------------------------------------------
# Tolerance helpers exported alongside Qreg
# ---------------------------------------------------------------------------


def test_amp_tol_complex128() -> None:
    assert amp_tol_for(torch.complex128) == AMP_TOL_C128


def test_amp_tol_complex64() -> None:
    assert amp_tol_for(torch.complex64) == AMP_TOL_C64


def test_amp_tol_unsupported_dtype_raises() -> None:
    with pytest.raises(ValueError, match=r"^qubit: amp_tol_for: "):
        amp_tol_for(torch.float64)


def test_prob_tol_complex128() -> None:
    assert prob_tol_for(torch.complex128) == PROB_TOL_C128


def test_prob_tol_complex64() -> None:
    assert prob_tol_for(torch.complex64) == PROB_TOL_C64


# ---------------------------------------------------------------------------
# Cross-device: construction on every available device
# ---------------------------------------------------------------------------


def test_construct_on_available_device(device: str) -> None:
    # The `device` fixture parametrises over whatever PyTorch can reach
    # on this host (always includes "cpu"; "mps" on Apple Silicon, etc.).
    # Construction must succeed on every available device with the
    # auto-selected dtype.
    q = Qreg(3, device=device)
    assert q.device.type == device
    # Default dtype must match policy.
    if device == "mps":
        assert q.dtype == torch.complex64
    else:
        assert q.dtype == torch.complex128


def test_init_basis_on_available_device(device: str) -> None:
    # The init_basis path must work end-to-end on every device, not just
    # CPU. This catches device-specific subtleties (e.g., the single-cell
    # assignment _amp[basis] = 1+0j has to be supported on the target).
    q = Qreg(3, device=device)
    q.init_basis(4)
    assert q.prob_of(4) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


def test_amplitudes_copy_always_cpu(device: str) -> None:
    # Regardless of the source device, amplitudes_copy must yield CPU.
    q = Qreg(3, device=device)
    q.init_basis(0)
    assert q.amplitudes_copy().device.type == "cpu"


================================================================================
FILE: implementation/python/tests/test_standart.py
================================================================================
"""Tests for the ``qubit.standart`` arithmetic helpers.

These mirror the corresponding ``/c`` and ``/go`` tests for the shared
function set. The Shor-relevant continued-fraction cases (period of 7
mod 15, period of 2 mod 21) are exercised in their actual readout form
``(c, 2**t, N)`` so a future Shor implementation can drop these helpers
in without rewriting the test contracts.
"""

from __future__ import annotations

import pytest

from qubit import (
    continued_fraction,
    gcd_u64,
    ilog2_u32,
    is_power_of_two,
    mod_pow,
)

# ---------------------------------------------------------------------------
# gcd_u64
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "a,b,want",
    [
        (0, 5, 5),
        (5, 0, 5),
        (12, 8, 4),
        (7, 13, 1),
        (15, 21, 3),
        (100, 75, 25),
        (0, 0, 0),
        # Beyond uint64: Python ints don't overflow, so gcd_u64 works on
        # arbitrarily large inputs. The function name still says "u64"
        # for parity with /c+/go, but the contract is "non-negative int".
        (10**20, 10**20 - 1, 1),
    ],
)
def test_gcd_basic(a: int, b: int, want: int) -> None:
    assert gcd_u64(a, b) == want


def test_gcd_rejects_negative_a() -> None:
    with pytest.raises(ValueError, match=r"^qubit: gcd_u64: a=-1"):
        gcd_u64(-1, 5)


def test_gcd_rejects_negative_b() -> None:
    with pytest.raises(ValueError, match=r"^qubit: gcd_u64: b=-3"):
        gcd_u64(7, -3)


def test_gcd_rejects_both_negative() -> None:
    # Either error is acceptable as long as one fires; the validation
    # is order-sensitive (a before b), so we expect the a-side message.
    with pytest.raises(ValueError, match=r"^qubit: gcd_u64: a=-5"):
        gcd_u64(-5, -7)


# ---------------------------------------------------------------------------
# mod_pow
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "base,exp,mod,want",
    [
        # Basic cases from the /go suite.
        (2, 10, 1000, 24),  # 1024 mod 1000
        (7, 0, 15, 1),
        (0, 0, 7, 1),  # 0**0 = 1 convention
        # Large exp -- pow handles square-and-multiply automatically.
        (2, 1 << 10, 1000000007, 812734592),
        # Mod over uint32.
        ((1 << 32) + 1, 5, (1 << 33) - 1, 5100273671),
        # mod == 1 -- everything reduces to 0.
        (5, 3, 1, 0),
        # Negative base (Python normalises): -2**3 mod 5 = -8 mod 5 = 2.
        (-2, 3, 5, 2),
    ],
)
def test_mod_pow_basic(base: int, exp: int, mod: int, want: int) -> None:
    assert mod_pow(base, exp, mod) == want


def test_mod_pow_rejects_negative_exp() -> None:
    with pytest.raises(ValueError, match=r"^qubit: mod_pow: exp=-1"):
        mod_pow(2, -1, 5)


def test_mod_pow_rejects_zero_mod() -> None:
    # mod == 0 would mean "anything mod 0", which Python's pow raises
    # ZeroDivisionError for. We surface this with the qubit: prefix.
    with pytest.raises(ValueError, match=r"^qubit: mod_pow: mod=0"):
        mod_pow(2, 5, 0)


def test_mod_pow_rejects_negative_mod() -> None:
    with pytest.raises(ValueError, match=r"^qubit: mod_pow: mod=-3"):
        mod_pow(2, 5, -3)


# ---------------------------------------------------------------------------
# continued_fraction
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "x,q,max_denom,want",
    [
        # 3/8 with a generous max-denom -- the exact denominator survives.
        (3, 8, 100, 8),
        # 1/3 with max-denom >= 3.
        (1, 3, 10, 3),
        # x == 0: the trivial 0/1 approximation, matching /go's x<=0 hatch.
        (0, 256, 15, 1),
        (0, 2048, 21, 1),
        # The fraction reduces internally: 2/4 is identical to 1/2 inside
        # Fraction(), so the returned denominator is the reduced one.
        (2, 4, 100, 2),
        # max_denom exactly equal to the fraction's reduced denominator.
        (3, 7, 7, 7),
    ],
)
def test_continued_fraction_basic(
    x: int, q: int, max_denom: int, want: int
) -> None:
    assert continued_fraction(x, q, max_denom) == want


@pytest.mark.parametrize(
    "c,want",
    [
        # /go's test_shor_period_a7_mod15 uses t=8 counting + n=4 target;
        # the true period of 7 mod 15 is 4, so c is one of {0, 64, 128,
        # 192} after the QFT. continued_fraction recovers a divisor of 4.
        (0, 1),
        (64, 4),
        (128, 2),
        (192, 4),
    ],
)
def test_continued_fraction_shor_a7_mod15(c: int, want: int) -> None:
    # Shor's period finder: c is the measured counting register, q is
    # 2**t (t=8 here), max_denominator is N (=15). The returned
    # denominator is the candidate period r.
    assert continued_fraction(c, 256, 15) == want


@pytest.mark.parametrize(
    "c,want",
    [
        # /go's test_shor_period_a2_mod21: t=11, true period r=6.
        # The QFT readout c is approximately k * 2048 / 6 for some k;
        # continued_fraction should recover a divisor of 6.
        (0, 1),
        (341, 6),   # 1/6
        (683, 3),   # 1/3
        (1024, 2),  # 1/2
        (1365, 3),  # 2/3
        (1707, 6),  # 5/6
    ],
)
def test_continued_fraction_shor_a2_mod21(c: int, want: int) -> None:
    assert continued_fraction(c, 2048, 21) == want


def test_continued_fraction_irrational_via_rational() -> None:
    # /go tests continued_fraction(sqrt(2), 10) -> denominator 5
    # (convergent 7/5). Python's signature doesn't take floats, but
    # we can express sqrt(2) ~ 14142135 / 10000000 and recover 7/5
    # via limit_denominator. Confirms the Stern-Brocot search matches
    # /go's hand-rolled convergent loop on a non-trivial irrational.
    assert continued_fraction(14142135, 10000000, 10) == 5


def test_continued_fraction_rejects_zero_q() -> None:
    with pytest.raises(ValueError, match=r"^qubit: continued_fraction: q=0"):
        continued_fraction(3, 0, 100)


def test_continued_fraction_rejects_negative_q() -> None:
    with pytest.raises(ValueError, match=r"^qubit: continued_fraction: q=-8"):
        continued_fraction(3, -8, 100)


def test_continued_fraction_rejects_zero_max_denominator() -> None:
    with pytest.raises(
        ValueError, match=r"^qubit: continued_fraction: max_denominator=0"
    ):
        continued_fraction(3, 8, 0)


def test_continued_fraction_rejects_negative_max_denominator() -> None:
    with pytest.raises(
        ValueError, match=r"^qubit: continued_fraction: max_denominator=-5"
    ):
        continued_fraction(3, 8, -5)


# ---------------------------------------------------------------------------
# is_power_of_two
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "x,want",
    [
        (1, True),
        (2, True),
        (4, True),
        (8, True),
        (1024, True),
        (1 << 30, True),
        (1 << 60, True),  # Python ints are unbounded; still detects.
        (0, False),
        (3, False),
        (6, False),
        (-1, False),
        (-2, False),
        (-1024, False),
    ],
)
def test_is_power_of_two(x: int, want: bool) -> None:
    assert is_power_of_two(x) is want


# ---------------------------------------------------------------------------
# ilog2_u32
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "x,want",
    [
        (1, 0),
        (2, 1),
        (3, 1),  # floor(log2(3)) = 1
        (4, 2),
        (7, 2),  # floor(log2(7)) = 2
        (8, 3),
        (1024, 10),
        (1 << 30, 30),
        (1 << 62, 62),  # Python ints have no width limit; the name is just parity.
    ],
)
def test_ilog2(x: int, want: int) -> None:
    assert ilog2_u32(x) == want


def test_ilog2_rejects_zero() -> None:
    with pytest.raises(ValueError, match=r"^qubit: ilog2_u32: x=0"):
        ilog2_u32(0)


def test_ilog2_rejects_negative() -> None:
    with pytest.raises(ValueError, match=r"^qubit: ilog2_u32: x=-3"):
        ilog2_u32(-3)


# ---------------------------------------------------------------------------
# Public re-export surface
# ---------------------------------------------------------------------------


def test_all_helpers_exported_from_qubit() -> None:
    # Sanity: every standart helper must be importable from the top-level
    # qubit package, not only from qubit.standart.
    import qubit

    for name in (
        "gcd_u64",
        "mod_pow",
        "continued_fraction",
        "is_power_of_two",
        "ilog2_u32",
    ):
        assert hasattr(qubit, name), f"qubit.{name} missing from re-exports"
        assert name in qubit.__all__, f"qubit.{name} missing from __all__"


================================================================================
FILE: implementation/python/tests/test_gates_single.py
================================================================================
"""Tests for the single-qubit gate primitives.

Covers ``apply_u`` (the workhorse) plus every named gate
(``apply_h`` / ``apply_x`` / ``apply_y`` / ``apply_z`` / ``apply_s`` /
``apply_t`` / ``apply_phase`` / ``apply_rx`` / ``apply_ry`` /
``apply_rz``), in both function-style ``apply_h(q, 0)`` and method-style
``q.apply_h(0)`` invocations.

The LSB-first qubit-to-axis mapping is the single most likely place to
introduce a silent endian flip; ``test_apply_x_lsb_axis_mapping`` and
``test_apply_x_msb_axis_mapping`` lock the convention in alongside
``tests/test_axis.py``.
"""

from __future__ import annotations

import cmath
import math

import pytest
import torch

from qubit import (
    Qreg,
    amp_tol_for,
    apply_h,
    apply_phase,
    apply_rx,
    apply_ry,
    apply_rz,
    apply_s,
    apply_t,
    apply_u,
    apply_x,
    apply_y,
    apply_z,
    prob_tol_for,
)

# ---------------------------------------------------------------------------
# Helper assertions
# ---------------------------------------------------------------------------


def _approx(q: Qreg, expected: complex, **_: object) -> object:
    """Return ``pytest.approx(expected)`` with the right dtype tolerance."""
    return pytest.approx(expected, abs=amp_tol_for(q.dtype))


# Convenient complex literal
INV_SQRT2 = 1.0 / math.sqrt(2.0)


# ---------------------------------------------------------------------------
# apply_u workhorse: shape, dtype, device, and target validation
# ---------------------------------------------------------------------------


def _eye2(q: Qreg) -> torch.Tensor:
    """Build a 2x2 identity tensor on q's device/dtype."""
    return torch.eye(2, dtype=q.dtype, device=q.device)


def test_apply_u_identity_is_noop() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(5)  # |101>
    apply_u(q, 1, _eye2(q))
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    for i in range(8):
        if i != 5:
            assert q.amplitude(i) == _approx(q, 0 + 0j)


def test_apply_u_with_explicit_x_matches_apply_x() -> None:
    # Custom u = X; result must match the named apply_x exactly.
    q1 = Qreg(2, device="cpu")
    q2 = Qreg(2, device="cpu")
    q1.init_basis(0)
    q2.init_basis(0)
    x = torch.tensor(
        [[0 + 0j, 1 + 0j],
         [1 + 0j, 0 + 0j]],
        dtype=q1.dtype, device=q1.device,
    )
    apply_u(q1, 0, x)
    apply_x(q2, 0)
    for i in range(4):
        assert q1.amplitude(i) == _approx(q1, q2.amplitude(i))


def test_apply_u_rejects_out_of_range_target() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: qubit_axis: target=5"):
        apply_u(q, 5, _eye2(q))


def test_apply_u_rejects_negative_target() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: qubit_axis: target=-1"):
        apply_u(q, -1, _eye2(q))


def test_apply_u_rejects_non_tensor() -> None:
    q = Qreg(2, device="cpu")
    bad = [[1.0, 0.0], [0.0, 1.0]]  # a plain Python list
    # No `# type: ignore` here: our mypy config follow-skips torch stubs,
    # so torch.Tensor resolves as Any and the type checker can't catch
    # this bad call statically. The runtime guard is the only defence.
    with pytest.raises(
        TypeError, match=r"^qubit: apply_u: matrix must be a torch.Tensor"
    ):
        apply_u(q, 0, bad)


def test_apply_u_rejects_wrong_shape() -> None:
    q = Qreg(2, device="cpu")
    wrong_shape = torch.eye(3, dtype=q.dtype, device=q.device)
    with pytest.raises(
        ValueError, match=r"^qubit: apply_u: matrix shape must be \(2, 2\)"
    ):
        apply_u(q, 0, wrong_shape)


def test_apply_u_rejects_wrong_dtype() -> None:
    q = Qreg(2, device="cpu")  # complex128
    wrong_dtype = torch.eye(2, dtype=torch.complex64, device=q.device)
    with pytest.raises(
        ValueError, match=r"^qubit: apply_u: matrix dtype must be"
    ):
        apply_u(q, 0, wrong_dtype)


@pytest.mark.skipif(
    not torch.backends.mps.is_available(),
    reason="device-mismatch test needs a second device; only MPS is locally available",
)
def test_apply_u_rejects_wrong_device() -> None:
    # Qreg on MPS, matrix on CPU -- caller bug; must raise.
    q = Qreg(2, device="mps")  # complex64
    cpu_matrix = torch.eye(2, dtype=q.dtype, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_u: matrix device must be"
    ):
        apply_u(q, 0, cpu_matrix)


# ---------------------------------------------------------------------------
# Hadamard
# ---------------------------------------------------------------------------


def test_apply_h_on_zero_produces_plus_state() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    assert q.amplitude(0) == _approx(q, INV_SQRT2 + 0j)
    assert q.amplitude(1) == _approx(q, INV_SQRT2 + 0j)


def test_apply_h_on_one_produces_minus_state() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_h(q, 0)
    # H|1> = (|0> - |1>) / sqrt(2)
    assert q.amplitude(0) == _approx(q, INV_SQRT2 + 0j)
    assert q.amplitude(1) == _approx(q, -INV_SQRT2 + 0j)


def test_apply_h_twice_is_identity() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(5)  # |101>
    apply_h(q, 1)
    apply_h(q, 1)
    for i in range(8):
        expected = 1 + 0j if i == 5 else 0 + 0j
        assert q.amplitude(i) == _approx(q, expected)


def test_apply_h_preserves_norm() -> None:
    q = Qreg(4, device="cpu")
    q.init_basis(7)
    apply_h(q, 0)
    apply_h(q, 1)
    apply_h(q, 2)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ---------------------------------------------------------------------------
# Pauli-X: bit flip + LSB axis check
# ---------------------------------------------------------------------------


def test_apply_x_on_zero_gives_one() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_x(q, 0)
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, 1 + 0j)


def test_apply_x_on_one_gives_zero() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_x(q, 0)
    assert q.amplitude(0) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_x_lsb_axis_mapping() -> None:
    # |001> = amp[1]; X on qubit 0 (LSB) -> |000> = amp[0].
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_x(q, 0)
    assert q.amplitude(0) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_x_msb_axis_mapping() -> None:
    # |001> = amp[1]; X on qubit 2 (MSB on 3 qubits) -> |101> = amp[5].
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_x(q, 2)
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_x_middle_qubit() -> None:
    # |001> = amp[1]; X on qubit 1 -> |011> = amp[3].
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_x(q, 1)
    assert q.amplitude(3) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


# ---------------------------------------------------------------------------
# Pauli-Y
# ---------------------------------------------------------------------------


def test_apply_y_on_zero_gives_i_times_one() -> None:
    # Y|0> = i|1>
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_y(q, 0)
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 1j)


def test_apply_y_on_one_gives_negative_i_times_zero() -> None:
    # Y|1> = -i|0>
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_y(q, 0)
    assert q.amplitude(0) == _approx(q, 0 - 1j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


# ---------------------------------------------------------------------------
# Pauli-Z
# ---------------------------------------------------------------------------


def test_apply_z_on_zero_is_identity() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_z(q, 0)
    assert q.amplitude(0) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_z_on_one_negates() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_z(q, 0)
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, -1 + 0j)


# ---------------------------------------------------------------------------
# S and T: S**2 == Z, T**4 == Z
# ---------------------------------------------------------------------------


def test_apply_s_squared_equals_z_on_one() -> None:
    # S|1> = i|1>; S(S|1>) = i*i|1> = -|1>. Same as Z|1>.
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_s(q, 0)
    apply_s(q, 0)
    assert q.amplitude(1) == _approx(q, -1 + 0j)


def test_apply_t_fourth_power_equals_z_on_one() -> None:
    # T|1> = e^(i pi/4)|1>; T^4|1> = e^(i pi)|1> = -|1>.
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    for _ in range(4):
        apply_t(q, 0)
    assert q.amplitude(1) == _approx(q, -1 + 0j)


# ---------------------------------------------------------------------------
# General phase gate
# ---------------------------------------------------------------------------


def test_apply_phase_pi_equals_z_on_one() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_phase(q, 0, math.pi)
    assert q.amplitude(1) == _approx(q, -1 + 0j)


def test_apply_phase_zero_is_identity() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_phase(q, 0, 0.0)
    assert q.amplitude(1) == _approx(q, 1 + 0j)


def test_apply_phase_arbitrary_angle() -> None:
    # Phase(pi/3)|1> = e^(i pi/3)|1>.
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_phase(q, 0, math.pi / 3.0)
    expected = cmath.exp(complex(0, math.pi / 3.0))
    assert q.amplitude(1) == _approx(q, expected)


# ---------------------------------------------------------------------------
# Rotations: RX, RY, RZ
# ---------------------------------------------------------------------------


def test_apply_rx_2pi_on_zero_is_minus_one() -> None:
    # RX(2 pi)|0> = -|0>. Classic sign flip after a full rotation.
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_rx(q, 0, 2 * math.pi)
    assert q.amplitude(0) == _approx(q, -1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_rx_pi_on_zero_gives_minus_i_one() -> None:
    # RX(pi)|0> = -i|1>: cos(pi/2)=0, -i sin(pi/2) = -i.
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_rx(q, 0, math.pi)
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, 0 - 1j)


def test_apply_ry_4pi_on_zero_is_identity() -> None:
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_ry(q, 0, 4 * math.pi)
    assert q.amplitude(0) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_ry_pi_on_zero_gives_one() -> None:
    # RY(pi)|0> = cos(pi/2)|0> + sin(pi/2)|1> = |1>.
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_ry(q, 0, math.pi)
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, 1 + 0j)


def test_apply_rz_diagonal_phase_on_zero() -> None:
    # RZ(theta)|0> = e^(-i theta/2)|0>. Test at theta = pi/2.
    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_rz(q, 0, math.pi / 2.0)
    expected = cmath.exp(complex(0, -math.pi / 4.0))
    assert q.amplitude(0) == _approx(q, expected)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_rz_diagonal_phase_on_one() -> None:
    # RZ(theta)|1> = e^(+i theta/2)|1>.
    q = Qreg(1, device="cpu")
    q.init_basis(1)
    apply_rz(q, 0, math.pi / 2.0)
    expected = cmath.exp(complex(0, math.pi / 4.0))
    assert q.amplitude(0) == _approx(q, 0 + 0j)
    assert q.amplitude(1) == _approx(q, expected)


# ---------------------------------------------------------------------------
# Method-style API: q.apply_h(0) must agree with apply_h(q, 0)
# ---------------------------------------------------------------------------


def test_method_apply_h_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(0)
    qm.init_basis(0)
    apply_h(qf, 0)
    qm.apply_h(0)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_x_matches_function() -> None:
    qf = Qreg(3, device="cpu")
    qm = Qreg(3, device="cpu")
    qf.init_basis(1)  # |001>
    qm.init_basis(1)
    apply_x(qf, 0)
    qm.apply_x(0)
    for i in range(8):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_u_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(0)
    qm.init_basis(0)
    u = torch.tensor(
        [[INV_SQRT2 + 0j, INV_SQRT2 + 0j],
         [INV_SQRT2 + 0j, -INV_SQRT2 + 0j]],
        dtype=qf.dtype, device=qf.device,
    )
    apply_u(qf, 0, u)
    # Build a fresh tensor for the method-style call too (otherwise the
    # method-style would alias the same tensor; same result, but the
    # test is sharper when each path uses its own matrix).
    u2 = torch.tensor(
        [[INV_SQRT2 + 0j, INV_SQRT2 + 0j],
         [INV_SQRT2 + 0j, -INV_SQRT2 + 0j]],
        dtype=qm.dtype, device=qm.device,
    )
    qm.apply_u(0, u2)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_rz_matches_function() -> None:
    qf = Qreg(1, device="cpu")
    qm = Qreg(1, device="cpu")
    qf.init_basis(0)
    qm.init_basis(0)
    apply_rz(qf, 0, 0.7)
    qm.apply_rz(0, 0.7)
    for i in range(2):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


# ---------------------------------------------------------------------------
# Norm preservation: every named gate is unitary, so norm stays at 1.0
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "gate",
    ["h", "x", "y", "z", "s", "t"],
)
def test_named_gate_preserves_norm_on_basis_state(gate: str) -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(3)
    {
        "h": apply_h,
        "x": apply_x,
        "y": apply_y,
        "z": apply_z,
        "s": apply_s,
        "t": apply_t,
    }[gate](q, 1)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


@pytest.mark.parametrize("theta", [0.0, 0.5, math.pi / 3.0, math.pi, 2.5])
def test_rotations_preserve_norm(theta: float) -> None:
    for apply_r in (apply_rx, apply_ry, apply_rz):
        q = Qreg(2, device="cpu")
        # Use a superposition starting state so the rotation actually
        # exercises both components, not just a basis-state phase.
        q.init_basis(0)
        apply_h(q, 0)
        apply_r(q, 0, theta)
        assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ---------------------------------------------------------------------------
# Device-parametrised sanity check: gates work on every available device
# ---------------------------------------------------------------------------


def test_apply_h_on_zero_across_devices(device: str) -> None:
    q = Qreg(1, device=device)
    q.init_basis(0)
    apply_h(q, 0)
    assert q.amplitude(0) == pytest.approx(
        INV_SQRT2 + 0j, abs=amp_tol_for(q.dtype)
    )
    assert q.amplitude(1) == pytest.approx(
        INV_SQRT2 + 0j, abs=amp_tol_for(q.dtype)
    )


def test_apply_x_lsb_across_devices(device: str) -> None:
    # Repeats the LSB-axis test on every available device to catch any
    # device-specific tensordot/movedim/reshape interaction.
    q = Qreg(3, device=device)
    q.init_basis(1)
    apply_x(q, 0)
    assert q.prob_of(0) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.prob_of(1) == pytest.approx(0.0, abs=prob_tol_for(q.dtype))


def test_apply_u_custom_matrix_across_devices(device: str) -> None:
    # Builds the matrix on the same device, then applies it. This is the
    # contract: caller pre-places the matrix; gate doesn't auto-transfer.
    q = Qreg(2, device=device)
    q.init_basis(0)
    x = torch.tensor(
        [[0 + 0j, 1 + 0j],
         [1 + 0j, 0 + 0j]],
        dtype=q.dtype, device=q.device,
    )
    apply_u(q, 0, x)
    assert q.prob_of(1) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.prob_of(0) == pytest.approx(0.0, abs=prob_tol_for(q.dtype))


================================================================================
FILE: implementation/python/tests/test_gates_controlled.py
================================================================================
"""Tests for the two-qubit controlled gates.

Covers ``apply_cu`` (the workhorse) plus the named controlled gates
(``apply_cnot`` / ``apply_cz`` / ``apply_controlled_phase``) and
``apply_swap``. The two highest-signal tests are:

* the Bell state ``H(0) -> CNOT(0, 1)`` producing
  ``(|00> + |11>) / sqrt(2)`` -- locks in the 4x4 layout convention and
  the permute / matmul / inverse-permute sequence at once;
* the non-adjacent ``CNOT(0, 2)`` on a 3-qubit register -- locks in the
  inverse-permutation accounting that adjacency would mask.
"""

from __future__ import annotations

import cmath
import math

import pytest
import torch

from qubit import (
    Qreg,
    amp_tol_for,
    apply_cnot,
    apply_controlled_phase,
    apply_cu,
    apply_cz,
    apply_h,
    apply_swap,
    prob_tol_for,
)

INV_SQRT2 = 1.0 / math.sqrt(2.0)


def _approx(q: Qreg, expected: complex) -> object:
    return pytest.approx(expected, abs=amp_tol_for(q.dtype))


# ---------------------------------------------------------------------------
# apply_cu workhorse
# ---------------------------------------------------------------------------


def _x_matrix(q: Qreg) -> torch.Tensor:
    return torch.tensor(
        [[0 + 0j, 1 + 0j], [1 + 0j, 0 + 0j]],
        dtype=q.dtype, device=q.device,
    )


def _identity_matrix(q: Qreg) -> torch.Tensor:
    return torch.eye(2, dtype=q.dtype, device=q.device)


def test_apply_cu_identity_is_noop() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(5)  # |101>
    apply_cu(q, 0, 1, _identity_matrix(q))
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    for i in range(8):
        if i != 5:
            assert q.amplitude(i) == _approx(q, 0 + 0j)


def test_apply_cu_custom_x_matches_apply_cnot() -> None:
    qf = Qreg(3, device="cpu")
    qm = Qreg(3, device="cpu")
    qf.init_basis(1)
    qm.init_basis(1)
    apply_cu(qf, 0, 1, _x_matrix(qf))
    apply_cnot(qm, 0, 1)
    for i in range(8):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_apply_cu_rejects_out_of_range_control() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_cu: control=5"):
        apply_cu(q, 5, 0, _identity_matrix(q))


def test_apply_cu_rejects_out_of_range_target() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_cu: target=9"):
        apply_cu(q, 0, 9, _identity_matrix(q))


def test_apply_cu_rejects_equal_control_target() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_cu: control == target == 1"):
        apply_cu(q, 1, 1, _identity_matrix(q))


def test_apply_cu_rejects_non_tensor() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        TypeError, match=r"^qubit: apply_cu: matrix must be a torch.Tensor"
    ):
        apply_cu(q, 0, 1, [[1, 0], [0, 1]])


def test_apply_cu_rejects_wrong_shape() -> None:
    q = Qreg(2, device="cpu")
    wrong = torch.eye(3, dtype=q.dtype, device=q.device)
    with pytest.raises(
        ValueError, match=r"^qubit: apply_cu: matrix shape must be \(2, 2\)"
    ):
        apply_cu(q, 0, 1, wrong)


def test_apply_cu_rejects_wrong_dtype() -> None:
    q = Qreg(2, device="cpu")  # complex128
    wrong = torch.eye(2, dtype=torch.complex64, device=q.device)
    with pytest.raises(
        ValueError, match=r"^qubit: apply_cu: matrix dtype must be"
    ):
        apply_cu(q, 0, 1, wrong)


# ---------------------------------------------------------------------------
# CNOT: bit flip, LSB / MSB axis, control=0 no-op, Bell state
# ---------------------------------------------------------------------------


def test_apply_cnot_control_zero_is_noop() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(2)  # |10> = q1=1, q0=0
    apply_cnot(q, 0, 1)  # control q0=0, target q1; no flip
    assert q.amplitude(2) == _approx(q, 1 + 0j)


def test_apply_cnot_lsb_mapping_flips_target() -> None:
    # |01> = amp[1] with control q0=1; flip target q1 -> |11> = amp[3].
    q = Qreg(2, device="cpu")
    q.init_basis(1)
    apply_cnot(q, 0, 1)
    assert q.amplitude(3) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_cnot_non_adjacent_control_target() -> None:
    # 3 qubits, control=0, target=2. |001> = amp[1]; control bit q0=1,
    # flip target q2 -> |101> = amp[5].
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_cnot(q, 0, 2)
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_cnot_msb_control_lsb_target() -> None:
    # 3 qubits, control=2 (MSB), target=0 (LSB). |100> = amp[4];
    # control q2=1, flip target q0 -> |101> = amp[5].
    q = Qreg(3, device="cpu")
    q.init_basis(4)
    apply_cnot(q, 2, 0)
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    assert q.amplitude(4) == _approx(q, 0 + 0j)


def test_bell_state_from_h_then_cnot() -> None:
    # The canonical 4x4-layout-and-permutation acceptance test.
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    apply_cnot(q, 0, 1)
    assert q.amplitude(0) == _approx(q, INV_SQRT2 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)
    assert q.amplitude(2) == _approx(q, 0 + 0j)
    assert q.amplitude(3) == _approx(q, INV_SQRT2 + 0j)


def test_cnot_preserves_norm() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    apply_h(q, 1)
    apply_cnot(q, 0, 2)
    apply_cnot(q, 1, 2)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ---------------------------------------------------------------------------
# Controlled-Z and controlled-Phase
# ---------------------------------------------------------------------------


def test_apply_cz_flips_only_eleven() -> None:
    # CZ adds a sign to |11> and leaves the other three computational
    # basis states alone. Build (|0> + |1>)^2 / 2 and check.
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    apply_h(q, 1)
    apply_cz(q, 0, 1)
    half = 0.5
    assert q.amplitude(0) == _approx(q, half + 0j)
    assert q.amplitude(1) == _approx(q, half + 0j)
    assert q.amplitude(2) == _approx(q, half + 0j)
    assert q.amplitude(3) == _approx(q, -half + 0j)


def test_apply_cz_on_eleven_basis_state_negates() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(3)  # |11>
    apply_cz(q, 0, 1)
    assert q.amplitude(3) == _approx(q, -1 + 0j)


def test_apply_controlled_phase_pi_equals_cz() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(3)
    qm.init_basis(3)
    apply_controlled_phase(qf, 0, 1, math.pi)
    apply_cz(qm, 0, 1)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_apply_controlled_phase_arbitrary_angle() -> None:
    # CPhase(theta)|11> = e^{i theta}|11>.
    q = Qreg(2, device="cpu")
    q.init_basis(3)
    apply_controlled_phase(q, 0, 1, math.pi / 5.0)
    expected = cmath.exp(complex(0, math.pi / 5.0))
    assert q.amplitude(3) == _approx(q, expected)


def test_apply_controlled_phase_no_effect_on_other_branches() -> None:
    # Only the |11> branch picks up the phase. Verify for the other three.
    for basis in (0, 1, 2):
        q = Qreg(2, device="cpu")
        q.init_basis(basis)
        apply_controlled_phase(q, 0, 1, math.pi / 3.0)
        assert q.amplitude(basis) == _approx(q, 1 + 0j)


# ---------------------------------------------------------------------------
# SWAP
# ---------------------------------------------------------------------------


def test_apply_swap_exchanges_01_and_10() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(1)  # |01>
    apply_swap(q, 0, 1)
    # |01> <-> |10>
    assert q.amplitude(2) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_apply_swap_leaves_00_and_11_alone() -> None:
    for basis in (0, 3):
        q = Qreg(2, device="cpu")
        q.init_basis(basis)
        apply_swap(q, 0, 1)
        assert q.amplitude(basis) == _approx(q, 1 + 0j)


def test_apply_swap_non_adjacent() -> None:
    # 3 qubits, swap LSB and MSB. |100> (amp[4]) -> |001> (amp[1]).
    q = Qreg(3, device="cpu")
    q.init_basis(4)
    apply_swap(q, 0, 2)
    assert q.amplitude(1) == _approx(q, 1 + 0j)
    assert q.amplitude(4) == _approx(q, 0 + 0j)


def test_apply_swap_rejects_equal() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_swap: a == b == 1"):
        apply_swap(q, 1, 1)


def test_apply_swap_rejects_out_of_range() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_swap: a=5"):
        apply_swap(q, 5, 0)
    with pytest.raises(ValueError, match=r"^qubit: apply_swap: b=8"):
        apply_swap(q, 0, 8)


# ---------------------------------------------------------------------------
# Method-style API
# ---------------------------------------------------------------------------


def test_method_apply_cnot_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(0)
    qm.init_basis(0)
    apply_h(qf, 0)
    qf_function_path = qf  # alias for clarity
    qm.apply_h(0)
    apply_cnot(qf_function_path, 0, 1)
    qm.apply_cnot(0, 1)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_cu_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(1)
    qm.init_basis(1)
    # Build the X matrix separately for each path to ensure they don't
    # alias the same tensor; identical contents though.
    apply_cu(qf, 0, 1, _x_matrix(qf))
    qm.apply_cu(0, 1, _x_matrix(qm))
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_swap_matches_function() -> None:
    qf = Qreg(3, device="cpu")
    qm = Qreg(3, device="cpu")
    qf.init_basis(1)
    qm.init_basis(1)
    apply_swap(qf, 0, 2)
    qm.apply_swap(0, 2)
    for i in range(8):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_cz_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(3)
    qm.init_basis(3)
    apply_cz(qf, 0, 1)
    qm.apply_cz(0, 1)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_controlled_phase_matches_function() -> None:
    qf = Qreg(2, device="cpu")
    qm = Qreg(2, device="cpu")
    qf.init_basis(3)
    qm.init_basis(3)
    apply_controlled_phase(qf, 0, 1, 0.7)
    qm.apply_controlled_phase(0, 1, 0.7)
    for i in range(4):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


# ---------------------------------------------------------------------------
# Device-parametrised sanity: Bell state and non-adjacent CNOT on every device
# ---------------------------------------------------------------------------


def test_bell_state_across_devices(device: str) -> None:
    q = Qreg(2, device=device)
    q.init_basis(0)
    apply_h(q, 0)
    apply_cnot(q, 0, 1)
    tol = amp_tol_for(q.dtype)
    assert q.amplitude(0) == pytest.approx(INV_SQRT2 + 0j, abs=tol)
    assert q.amplitude(3) == pytest.approx(INV_SQRT2 + 0j, abs=tol)
    assert q.amplitude(1) == pytest.approx(0 + 0j, abs=tol)
    assert q.amplitude(2) == pytest.approx(0 + 0j, abs=tol)


def test_non_adjacent_cnot_across_devices(device: str) -> None:
    # The non-adjacent case is the one most likely to expose a
    # permutation-accounting bug, so we run it on every device.
    q = Qreg(3, device=device)
    q.init_basis(1)
    apply_cnot(q, 0, 2)
    assert q.prob_of(5) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.prob_of(1) == pytest.approx(0.0, abs=prob_tol_for(q.dtype))


def test_swap_across_devices(device: str) -> None:
    q = Qreg(3, device=device)
    q.init_basis(4)  # |100>
    apply_swap(q, 0, 2)  # exchange LSB and MSB
    assert q.prob_of(1) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.prob_of(4) == pytest.approx(0.0, abs=prob_tol_for(q.dtype))


================================================================================
FILE: implementation/python/tests/test_gates_multi.py
================================================================================
"""Tests for the multi-controlled gate primitives.

Covers :func:`apply_multi_controlled_z` (Grover diffusion's phase flip)
and :func:`apply_multi_controlled_x` (generalised Toffoli). Both target
the highest-likelihood failure mode (the MCX target-bit=0 mask) and
exhaustive validation.
"""

from __future__ import annotations

import math

import pytest

from qubit import (
    Qreg,
    amp_tol_for,
    apply_h,
    apply_multi_controlled_x,
    apply_multi_controlled_z,
    prob_tol_for,
)


def _approx(q: Qreg, expected: complex) -> object:
    return pytest.approx(expected, abs=amp_tol_for(q.dtype))


# ===========================================================================
# apply_multi_controlled_z
# ===========================================================================


def test_mcz_three_controls_flips_only_all_ones_in_superposition() -> None:
    # Prepare (|0> + |1>)^3 / sqrt(8) and apply MCZ over the three
    # qubits. Only the |111> amplitude should be negated; the other
    # seven stay at +1/sqrt(8).
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    apply_multi_controlled_z(q, [0, 1, 2])
    inv8 = 1.0 / math.sqrt(8.0)
    for i in range(8):
        expected = -inv8 if i == 7 else inv8
        assert q.amplitude(i) == _approx(q, expected + 0j)


def test_mcz_two_controls_negates_only_q0_q1_both_one() -> None:
    # 3 qubits, MCZ on controls [0, 1]. Should negate amp[3] (|011>) and
    # amp[7] (|111>) -- both have q0=q1=1. Other six are untouched.
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    apply_multi_controlled_z(q, [0, 1])
    inv8 = 1.0 / math.sqrt(8.0)
    expected_signs = [+1, +1, +1, -1, +1, +1, +1, -1]
    for i in range(8):
        assert q.amplitude(i) == _approx(q, expected_signs[i] * inv8 + 0j)


def test_mcz_single_control_is_z_on_that_qubit() -> None:
    # With one control [c], MCZ degenerates to "negate any basis state
    # with bit c set" -- exactly apply_z(q, c) up to global phase.
    q = Qreg(2, device="cpu")
    q.init_basis(3)  # |11>
    apply_multi_controlled_z(q, [0])
    # Negates because q0 is 1.
    assert q.amplitude(3) == _approx(q, -1 + 0j)


def test_mcz_non_contiguous_controls() -> None:
    # 4 qubits, controls=[0, 3] (skipping qubits 1 and 2). Only
    # amplitudes with both q0=1 AND q3=1 should be negated.
    q = Qreg(4, device="cpu")
    q.init_basis(0)
    for i in range(4):
        apply_h(q, i)
    apply_multi_controlled_z(q, [0, 3])
    inv16 = 1.0 / 4.0  # 1/sqrt(16)
    for i in range(16):
        q0 = i & 1
        q3 = (i >> 3) & 1
        sign = -1 if (q0 == 1 and q3 == 1) else +1
        assert q.amplitude(i) == _approx(q, sign * inv16 + 0j)


def test_mcz_leaves_norm_unchanged() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    apply_multi_controlled_z(q, [0, 1, 2])
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


def test_mcz_rejects_empty_controls() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_multi_controlled_z: controls must be a non-empty",
    ):
        apply_multi_controlled_z(q, [])


def test_mcz_rejects_out_of_range_control() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_z: control=5"
    ):
        apply_multi_controlled_z(q, [0, 5])


def test_mcz_rejects_negative_control() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_z: control=-1"
    ):
        apply_multi_controlled_z(q, [0, -1])


def test_mcz_rejects_duplicate_controls() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_z: duplicate control=1"
    ):
        apply_multi_controlled_z(q, [1, 1])


def test_mcz_accepts_tuple_controls() -> None:
    # The signature is list[int] | tuple[int, ...]; both work.
    q = Qreg(2, device="cpu")
    q.init_basis(3)
    apply_multi_controlled_z(q, (0, 1))
    assert q.amplitude(3) == _approx(q, -1 + 0j)


# ===========================================================================
# apply_multi_controlled_x (generalised Toffoli)
# ===========================================================================


def test_mcx_toffoli_flips_target_when_both_controls_on() -> None:
    # |011> -> Toffoli with controls=[0,1], target=2 -> |111>.
    q = Qreg(3, device="cpu")
    q.init_basis(3)
    apply_multi_controlled_x(q, [0, 1], 2)
    assert q.amplitude(7) == _approx(q, 1 + 0j)
    assert q.amplitude(3) == _approx(q, 0 + 0j)


def test_mcx_no_flip_when_any_control_zero() -> None:
    # |001> has q0=1 but q1=0; Toffoli leaves it alone.
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_multi_controlled_x(q, [0, 1], 2)
    assert q.amplitude(1) == _approx(q, 1 + 0j)


def test_mcx_single_control_is_cnot() -> None:
    # Toffoli with one control degenerates to CNOT(control, target).
    # |001> with controls=[0], target=2 -> flip q2 -> |101>.
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_multi_controlled_x(q, [0], 2)
    assert q.amplitude(5) == _approx(q, 1 + 0j)
    assert q.amplitude(1) == _approx(q, 0 + 0j)


def test_mcx_inverse_is_self() -> None:
    # X^2 = I and the controls are a permutation, so MCX^2 must be I.
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    snapshot = q.amplitudes_copy()
    apply_multi_controlled_x(q, [0, 1], 2)
    apply_multi_controlled_x(q, [0, 1], 2)
    after = q.amplitudes_copy()
    for i in range(8):
        diff = abs(complex(snapshot[i].item()) - complex(after[i].item()))
        assert diff <= amp_tol_for(q.dtype), (
            f"MCX^2 != identity at amp[{i}]: diff={diff}"
        )


def test_mcx_three_controls_flips_only_when_all_on() -> None:
    # 4 qubits, controls=[0,1,2], target=3. |0111> = amp[7] has
    # q0=q1=q2=1 -> flip q3 -> |1111> = amp[15]. Other near-misses
    # (one control off) should be untouched.
    for basis, want in [
        (7, 15),    # all three controls on -> flip
        (3, 3),     # only q0,q1 on -> no flip
        (5, 5),     # only q0,q2 on -> no flip
        (6, 6),     # only q1,q2 on -> no flip
        (1, 1),     # only q0 on   -> no flip
    ]:
        q = Qreg(4, device="cpu")
        q.init_basis(basis)
        apply_multi_controlled_x(q, [0, 1, 2], 3)
        assert q.amplitude(want) == _approx(q, 1 + 0j), (
            f"MCX on basis={basis} -> want amp[{want}]=1, got {q.amplitude(want)}"
        )


def test_mcx_lsb_target_mapping() -> None:
    # Target the LSB: controls=[1, 2], target=0. |110> = amp[6];
    # q1=q2=1 -> flip q0 -> |111> = amp[7].
    q = Qreg(3, device="cpu")
    q.init_basis(6)
    apply_multi_controlled_x(q, [1, 2], 0)
    assert q.amplitude(7) == _approx(q, 1 + 0j)
    assert q.amplitude(6) == _approx(q, 0 + 0j)


def test_mcx_preserves_norm_under_superposition() -> None:
    q = Qreg(4, device="cpu")
    q.init_basis(0)
    for i in range(4):
        apply_h(q, i)
    apply_multi_controlled_x(q, [0, 1], 3)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


def test_mcx_rejects_empty_controls() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_multi_controlled_x: controls must be a non-empty",
    ):
        apply_multi_controlled_x(q, [], 0)


def test_mcx_rejects_out_of_range_target() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_x: target=5"
    ):
        apply_multi_controlled_x(q, [0], 5)


def test_mcx_rejects_out_of_range_control() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_x: control=7"
    ):
        apply_multi_controlled_x(q, [0, 7], 2)


def test_mcx_rejects_duplicate_controls() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_x: duplicate control=1"
    ):
        apply_multi_controlled_x(q, [1, 1], 2)


def test_mcx_rejects_target_in_controls() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_multi_controlled_x: control 2 == target"
    ):
        apply_multi_controlled_x(q, [0, 2], 2)


def test_mcx_accepts_tuple_controls() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(3)
    apply_multi_controlled_x(q, (0, 1), 2)
    assert q.amplitude(7) == _approx(q, 1 + 0j)


# ===========================================================================
# Method-style API
# ===========================================================================


def test_method_apply_mcz_matches_function() -> None:
    qf = Qreg(3, device="cpu")
    qm = Qreg(3, device="cpu")
    qf.init_basis(0)
    qm.init_basis(0)
    for i in range(3):
        apply_h(qf, i)
        qm.apply_h(i)
    apply_multi_controlled_z(qf, [0, 1, 2])
    qm.apply_multi_controlled_z([0, 1, 2])
    for i in range(8):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


def test_method_apply_mcx_matches_function() -> None:
    qf = Qreg(3, device="cpu")
    qm = Qreg(3, device="cpu")
    qf.init_basis(3)
    qm.init_basis(3)
    apply_multi_controlled_x(qf, [0, 1], 2)
    qm.apply_multi_controlled_x([0, 1], 2)
    for i in range(8):
        assert qf.amplitude(i) == _approx(qf, qm.amplitude(i))


# ===========================================================================
# Device-parametrised sanity
# ===========================================================================


def test_mcz_all_controls_across_devices(device: str) -> None:
    # Same uniform-superposition MCZ test on every available device.
    q = Qreg(3, device=device)
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    apply_multi_controlled_z(q, [0, 1, 2])
    inv8 = 1.0 / math.sqrt(8.0)
    tol = amp_tol_for(q.dtype)
    assert q.amplitude(7) == pytest.approx(-inv8 + 0j, abs=tol)
    assert q.amplitude(0) == pytest.approx(inv8 + 0j, abs=tol)


def test_mcx_toffoli_across_devices(device: str) -> None:
    q = Qreg(3, device=device)
    q.init_basis(3)
    apply_multi_controlled_x(q, [0, 1], 2)
    assert q.prob_of(7) == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    assert q.prob_of(3) == pytest.approx(0.0, abs=prob_tol_for(q.dtype))


================================================================================
FILE: implementation/python/tests/test_measure.py
================================================================================
"""Tests for measurement, sampling, clone, and dump.

Probabilistic tests use a pinned seed and the per-Qreg CPU RNG; the
sampling helpers are deterministic given a seed because measurement
sampling crosses to CPU at the readout boundary (see Phase 0+1's
design rationale).
"""

from __future__ import annotations

import math
from collections import Counter

import pytest
import torch

from qubit import (
    Qreg,
    amp_tol_for,
    apply_h,
    apply_x,
    clone,
    dump,
    measure_all,
    measure_qubit,
    prob_tol_for,
    sample_distribution,
)

INV_SQRT2 = 1.0 / math.sqrt(2.0)


# ===========================================================================
# measure_qubit -- deterministic on basis states
# ===========================================================================


def test_measure_qubit_basis_state_each_qubit() -> None:
    # |101> = amp[5]: q0=1, q1=0, q2=1. Measure each in turn; on a
    # basis state the outcome is fully determined.
    for target, want in [(0, 1), (1, 0), (2, 1)]:
        q = Qreg(3, device="cpu", seed=42)
        q.init_basis(5)
        outcome = measure_qubit(q, target)
        assert outcome == want, (
            f"measure_qubit(target={target}) on |101> = {outcome}, want {want}"
        )


def test_measure_qubit_collapse_preserves_norm() -> None:
    q = Qreg(2, device="cpu", seed=7)
    q.init_basis(0)
    apply_h(q, 0)
    measure_qubit(q, 0)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


def test_measure_qubit_collapse_zeros_other_branch() -> None:
    # After measuring qubit 0 of |+>, the state must be exactly a
    # basis state of qubit 0: either amp[0]=1, amp[1]=0 or vice versa.
    q = Qreg(1, device="cpu", seed=99)
    q.init_basis(0)
    apply_h(q, 0)
    outcome = measure_qubit(q, 0)
    tol = amp_tol_for(q.dtype)
    if outcome == 0:
        assert q.amplitude(0) == pytest.approx(1 + 0j, abs=tol)
        assert q.amplitude(1) == pytest.approx(0 + 0j, abs=tol)
    else:
        assert q.amplitude(0) == pytest.approx(0 + 0j, abs=tol)
        assert q.amplitude(1) == pytest.approx(1 + 0j, abs=tol)


def test_measure_qubit_distribution_on_plus_state() -> None:
    # H|0> -> measure qubit 0 in a loop; expect roughly 50/50 outcomes.
    # We re-prepare the +-state each iteration (cheap, no RNG cost), so
    # each measurement consumes one RNG draw from the same q._gen.
    n_shots = 1000
    q = Qreg(1, device="cpu", seed=2026)
    counts = [0, 0]
    for _ in range(n_shots):
        q.init_basis(0)
        apply_h(q, 0)
        counts[measure_qubit(q, 0)] += 1
    # 3-sigma window for binomial(1000, 0.5) is ~ 500 +/- 47. Use
    # +/-100 for safety; the test must not flake on slow CI.
    assert 400 < counts[0] < 600, (
        f"H|0> 1000 shots: {counts[0]} zeros (expected ~500)"
    )
    assert 400 < counts[1] < 600


def test_measure_qubit_rejects_out_of_range() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: measure_qubit: target=5"
    ):
        measure_qubit(q, 5)


def test_measure_qubit_rejects_negative() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: measure_qubit: target=-1"
    ):
        measure_qubit(q, -1)


def test_measure_qubit_raises_qubit_prefix_on_zero_norm() -> None:
    # A fresh Qreg is now in |0...0> (amp[0] = 1), so we explicitly
    # zero the amplitude vector here to exercise the zero-norm guard.
    # measure_qubit must raise a "qubit:"-prefixed RuntimeError rather
    # than crashing with a bare math error or producing an unbiased
    # result.
    q = Qreg(2, device="cpu")
    q._amp.zero_()
    with pytest.raises(
        RuntimeError, match=r"^qubit: measure_qubit: total probability"
    ):
        measure_qubit(q, 0)


def test_measure_qubit_handles_unnormalised_state() -> None:
    # Hand-scale H|0> by a factor of 2 so |amp|^2 sums to 4 rather than 1.
    # The sampling step must use p0/total (not the raw p0), and the
    # post-collapse state must still be unit-norm. Old code with the
    # un-normalised `u < p0` comparison would always pick outcome=0 here
    # because u in [0, 1) is always < p0=2.
    q = Qreg(1, device="cpu", seed=42)
    q.init_basis(0)
    apply_h(q, 0)
    q._amp.mul_(2.0)
    # Pre-condition: state is scaled to norm 4.
    assert q.norm() == pytest.approx(4.0, abs=prob_tol_for(q.dtype))
    outcome = measure_qubit(q, 0)
    assert outcome in (0, 1)
    # Post-condition: collapse to a basis state with unit norm.
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    tol = amp_tol_for(q.dtype)
    if outcome == 0:
        assert q.amplitude(0) == pytest.approx(1 + 0j, abs=tol)
        assert q.amplitude(1) == pytest.approx(0 + 0j, abs=tol)
    else:
        assert q.amplitude(0) == pytest.approx(0 + 0j, abs=tol)
        assert q.amplitude(1) == pytest.approx(1 + 0j, abs=tol)


# ===========================================================================
# measure_all
# ===========================================================================


def test_measure_all_basis_state_is_deterministic() -> None:
    q = Qreg(3, device="cpu", seed=11)
    q.init_basis(6)
    assert measure_all(q) == 6


def test_measure_all_collapses_to_chosen_basis() -> None:
    q = Qreg(2, device="cpu", seed=7)
    q.init_basis(0)
    apply_h(q, 0)
    apply_h(q, 1)
    chosen = measure_all(q)
    # After collapse, only amp[chosen] is non-zero.
    tol = amp_tol_for(q.dtype)
    for i in range(4):
        expected = 1 + 0j if i == chosen else 0 + 0j
        assert q.amplitude(i) == pytest.approx(expected, abs=tol)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


def test_measure_all_raises_qubit_prefix_on_zero_norm() -> None:
    # A fresh Qreg is now in |0...0> (amp[0] = 1), so we explicitly
    # zero the amplitude vector here to exercise the zero-norm guard.
    # Calling measure_all must raise a "qubit:"-prefixed RuntimeError
    # rather than leaking torch.multinomial's bare RuntimeError.
    q = Qreg(2, device="cpu")
    q._amp.zero_()
    with pytest.raises(
        RuntimeError, match=r"^qubit: measure_all: total probability"
    ):
        measure_all(q)


def test_measure_all_distribution_on_uniform_superposition() -> None:
    # Apply H on each qubit of a 3-qubit |0> -> uniform 1/sqrt(8) over
    # all 8 basis states. Use sample_distribution to count outcomes;
    # each bin should be ~125 out of 1000.
    q = Qreg(3, device="cpu", seed=2026)
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    outcomes = sample_distribution(q, 1000)
    counts = Counter(outcomes)
    # 3-sigma window for binomial(1000, 1/8) is 125 +/- ~33. Use a
    # generous +/-50 to avoid flakes.
    for k in range(8):
        assert 75 < counts[k] < 175, (
            f"uniform 1000 shots: amp[{k}] count = {counts[k]}"
        )


# ===========================================================================
# sample_distribution
# ===========================================================================


def test_sample_distribution_does_not_mutate_original() -> None:
    q = Qreg(2, device="cpu", seed=42)
    q.init_basis(0)
    apply_h(q, 0)  # |+> on qubit 0; uniform over {0, 1}
    snapshot_before = q.amplitudes_copy()
    sample_distribution(q, 50)
    snapshot_after = q.amplitudes_copy()
    # The full amplitude vector must be unchanged.
    assert torch.equal(snapshot_before, snapshot_after)


def test_sample_distribution_returns_list_of_ints() -> None:
    q = Qreg(2, device="cpu", seed=1)
    q.init_basis(3)
    out = sample_distribution(q, 5)
    assert isinstance(out, list)
    assert all(isinstance(x, int) for x in out)
    # |11> is deterministic; every shot must return 3.
    assert out == [3, 3, 3, 3, 3]


def test_sample_distribution_zero_shots_returns_empty() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    assert sample_distribution(q, 0) == []


def test_sample_distribution_rejects_negative_shots() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: sample_distribution: shots=-1"
    ):
        sample_distribution(q, -1)


def test_sample_distribution_preserves_norm() -> None:
    # The function must restore the original (norm=1) state, including
    # after sampling has internally collapsed the register.
    q = Qreg(3, device="cpu", seed=2026)
    q.init_basis(0)
    for i in range(3):
        apply_h(q, i)
    sample_distribution(q, 100)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ===========================================================================
# clone
# ===========================================================================


def test_clone_amp_is_independent() -> None:
    q1 = Qreg(2, device="cpu", seed=42)
    q1.init_basis(1)
    q2 = clone(q1)
    # Mutate the clone: apply X to flip the LSB of basis 1 -> basis 0.
    apply_x(q2, 0)
    # Original is unchanged.
    tol = amp_tol_for(q1.dtype)
    assert q1.amplitude(1) == pytest.approx(1 + 0j, abs=tol)
    assert q1.amplitude(0) == pytest.approx(0 + 0j, abs=tol)
    # Clone has the mutation.
    assert q2.amplitude(0) == pytest.approx(1 + 0j, abs=tol)
    assert q2.amplitude(1) == pytest.approx(0 + 0j, abs=tol)


def test_clone_amp_matches_initially() -> None:
    q1 = Qreg(3, device="cpu", seed=2026)
    q1.init_basis(0)
    for i in range(3):
        apply_h(q1, i)
    q2 = clone(q1)
    snap1 = q1.amplitudes_copy()
    snap2 = q2.amplitudes_copy()
    assert torch.equal(snap1, snap2)


def test_clone_preserves_metadata() -> None:
    q1 = Qreg(4, device="cpu", seed=1, dtype=torch.complex128)
    q2 = clone(q1)
    assert q2.n_qubits == q1.n_qubits
    assert q2.device == q1.device
    assert q2.dtype == q1.dtype


def test_clone_rng_state_starts_equal() -> None:
    # The clone copies q's RNG state. Until either advances, both
    # generators draw identical random sequences -- including the same
    # measurement outcome on identical superpositions.
    q1 = Qreg(1, device="cpu", seed=42)
    q1.init_basis(0)
    apply_h(q1, 0)
    q2 = clone(q1)
    # Both clones are |+> with the same RNG state. The first
    # measurement on each must give the same outcome.
    outcome1 = measure_qubit(q1, 0)
    outcome2 = measure_qubit(q2, 0)
    assert outcome1 == outcome2


def test_clone_rng_independence_after_first_draw() -> None:
    # After clone, one register advancing its RNG must not advance the
    # other's. Verify by drawing from each and checking they STILL
    # match (because both were drawing from independent generators
    # that happened to have identical states at the start).
    q1 = Qreg(1, device="cpu", seed=42)
    q2 = clone(q1)
    # Advance q1's generator three times.
    for _ in range(3):
        torch.rand((), generator=q1._gen)
    # q2's generator is still at the initial state. Draw a value from
    # each and check that q2's value is what q1's FIRST value was.
    # That requires resetting; cleaner test: verify q1._gen.get_state()
    # is different from q2._gen.get_state() after q1 advanced.
    state1 = q1._gen.get_state()
    state2 = q2._gen.get_state()
    assert not torch.equal(state1, state2), (
        "after advancing q1's RNG, the two generators' states must "
        "differ -- they're independent"
    )


# ===========================================================================
# dump
# ===========================================================================


def test_dump_on_basis_state() -> None:
    q = Qreg(3, device="cpu")
    q.init_basis(5)
    result = dump(q)
    assert result == [(5, 1 + 0j)]


def test_dump_on_superposition_returns_all_branches() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    result = dump(q)
    # Should have two entries: amp[0] and amp[1].
    assert len(result) == 2
    indices = sorted(idx for idx, _ in result)
    assert indices == [0, 1]
    for _, amp in result:
        assert abs(amp - (INV_SQRT2 + 0j)) <= amp_tol_for(q.dtype)


def test_dump_threshold_filters() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    # H|0> amplitudes are 1/sqrt(2) ~= 0.707. Threshold 0.8 filters all.
    assert dump(q, threshold=0.8) == []
    # Threshold 0.5 keeps both.
    assert len(dump(q, threshold=0.5)) == 2


def test_dump_on_zero_state_returns_empty() -> None:
    # A fresh Qreg is now in |0...0> (amp[0] = 1), so we explicitly
    # zero the amplitude vector here to exercise the empty-dump path.
    q = Qreg(3, device="cpu")
    q._amp.zero_()
    assert dump(q) == []


def test_dump_returns_tuples_of_int_and_complex() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(2)
    result = dump(q)
    assert len(result) == 1
    basis, amp = result[0]
    assert isinstance(basis, int)
    assert isinstance(amp, complex)
    assert basis == 2


def test_dump_rejects_negative_threshold() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    with pytest.raises(ValueError, match=r"^qubit: dump: threshold=-0.5"):
        dump(q, threshold=-0.5)


# ===========================================================================
# Method-style API: q.X() must match the function-style X(q)
# ===========================================================================


def test_method_measure_qubit_matches_function() -> None:
    # Both call paths share the same RNG state initially, so the
    # outcome must match.
    q1 = Qreg(2, device="cpu", seed=42)
    q1.init_basis(0)
    apply_h(q1, 0)
    q2 = clone(q1)
    out_function = measure_qubit(q1, 0)
    out_method = q2.measure_qubit(0)
    assert out_function == out_method


def test_method_measure_all_matches_function() -> None:
    q1 = Qreg(3, device="cpu", seed=2026)
    q1.init_basis(0)
    for i in range(3):
        apply_h(q1, i)
    q2 = clone(q1)
    out_function = measure_all(q1)
    out_method = q2.measure_all()
    assert out_function == out_method


def test_method_sample_distribution_matches_function() -> None:
    q1 = Qreg(2, device="cpu", seed=11)
    q1.init_basis(0)
    apply_h(q1, 0)
    q2 = clone(q1)
    out_function = sample_distribution(q1, 50)
    out_method = q2.sample_distribution(50)
    assert out_function == out_method


def test_method_dump_matches_function() -> None:
    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_h(q, 0)
    assert dump(q) == q.dump()


def test_method_clone_matches_function() -> None:
    q1 = Qreg(2, device="cpu", seed=42)
    q1.init_basis(1)
    c_function = clone(q1)
    c_method = q1.clone()
    # Both clones have the same amp and RNG state; their first draw
    # must match.
    assert torch.equal(c_function.amplitudes_copy(), c_method.amplitudes_copy())


# ===========================================================================
# Device-parametrised: basis measurement + dump on every available device
# ===========================================================================


def test_measure_qubit_basis_across_devices(device: str) -> None:
    q = Qreg(3, device=device, seed=42)
    q.init_basis(5)  # |101>
    # Deterministic measurement on every device.
    assert measure_qubit(q, 0) == 1


def test_measure_all_basis_across_devices(device: str) -> None:
    q = Qreg(3, device=device, seed=42)
    q.init_basis(6)
    assert measure_all(q) == 6


def test_dump_across_devices(device: str) -> None:
    q = Qreg(3, device=device)
    q.init_basis(5)
    result = dump(q)
    assert result == [(5, 1 + 0j)]


================================================================================
FILE: implementation/python/tests/test_qft.py
================================================================================
"""Tests for the Quantum Fourier Transform.

The five highest-signal tests:

* ``test_qft_on_zero_state_is_uniform`` -- QFT|0> = uniform 1/sqrt(N).
* ``test_qft_on_basis_state_matches_analytic`` -- locks the
  ``exp(2 pi i x y / N) / sqrt(N)`` formula across multiple
  ``(n, x)`` pairs and incidentally locks the bit-reversal-swap
  convention (without the final swaps the output would land at
  bit-reversed indices).
* ``test_qft_then_qft_inverse_round_trip`` -- exhaustive over basis
  states for n = 1..4, locks the inverse is a TRUE inverse.
* ``test_qft_inverse_then_qft_round_trip`` -- same, opposite order.
* ``test_qft_subregister_preserves_outer_qubits`` -- sub-register
  ``start=1, n=2`` on a 3-qubit register leaves the outer qubit
  coherent.
"""

from __future__ import annotations

import cmath
import math

import pytest

from qubit import (
    Qreg,
    amp_tol_for,
    apply_qft,
    apply_qft_inverse,
    prob_tol_for,
)

# ===========================================================================
# QFT on |0> = uniform superposition
# ===========================================================================


@pytest.mark.parametrize("n", [1, 2, 3, 4])
def test_qft_on_zero_state_is_uniform(n: int) -> None:
    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_qft(q)
    expected = 1.0 / math.sqrt(1 << n)
    tol = amp_tol_for(q.dtype)
    for i in range(1 << n):
        assert q.amplitude(i) == pytest.approx(expected + 0j, abs=tol), (
            f"QFT|0> n={n} amp[{i}] = {q.amplitude(i)}, expect {expected}"
        )


# ===========================================================================
# QFT on |x> matches the analytic formula
# ===========================================================================


@pytest.mark.parametrize(
    "n,x",
    [
        # n=2: every basis state
        (2, 0), (2, 1), (2, 2), (2, 3),
        # n=3: a representative sweep
        (3, 0), (3, 1), (3, 3), (3, 5), (3, 7),
        # n=4: a couple of cases (full sweep is 16 -- covered by
        # round-trip tests)
        (4, 1), (4, 11),
    ],
)
def test_qft_on_basis_state_matches_analytic(n: int, x: int) -> None:
    # amp[y] = exp(2 pi i * x * y / N) / sqrt(N) for input |x>.
    # The final bit-reversal swap is what puts the output in natural
    # binary order -- without it amp would land at bit_reverse(y).
    q = Qreg(n, device="cpu")
    q.init_basis(x)
    apply_qft(q)
    big_n = 1 << n
    tol = amp_tol_for(q.dtype)
    for y in range(big_n):
        expected = cmath.exp(complex(0, 2.0 * math.pi * x * y / big_n)) / math.sqrt(big_n)
        assert q.amplitude(y) == pytest.approx(expected, abs=tol), (
            f"QFT|{x}> n={n} amp[{y}]={q.amplitude(y)}, expect {expected}"
        )


# ===========================================================================
# Round-trip: QFT then inverse, and the other order
# ===========================================================================


@pytest.mark.parametrize("n", [1, 2, 3, 4])
def test_qft_then_qft_inverse_round_trip(n: int) -> None:
    # Every basis state x must map back to |x> after QFT followed by
    # inverse-QFT. Locks that the inverse is a TRUE inverse, not an
    # equivalent "QFT-three-times" (which would also work, since QFT^4
    # = I, but at three times the gate cost -- documented in qft.py).
    tol = amp_tol_for(Qreg(1, device="cpu").dtype)
    for x in range(1 << n):
        q = Qreg(n, device="cpu")
        q.init_basis(x)
        apply_qft(q)
        apply_qft_inverse(q)
        for i in range(1 << n):
            expected = 1 + 0j if i == x else 0 + 0j
            assert q.amplitude(i) == pytest.approx(expected, abs=tol), (
                f"round-trip n={n} basis={x} amp[{i}]={q.amplitude(i)}"
            )


@pytest.mark.parametrize("n", [1, 2, 3])
def test_qft_inverse_then_qft_round_trip(n: int) -> None:
    tol = amp_tol_for(Qreg(1, device="cpu").dtype)
    for x in range(1 << n):
        q = Qreg(n, device="cpu")
        q.init_basis(x)
        apply_qft_inverse(q)
        apply_qft(q)
        for i in range(1 << n):
            expected = 1 + 0j if i == x else 0 + 0j
            assert q.amplitude(i) == pytest.approx(expected, abs=tol), (
                f"inverse-then-forward n={n} basis={x} amp[{i}]={q.amplitude(i)}"
            )


def test_qft_preserves_norm() -> None:
    # Unitarity check: norm stays 1 through QFT, inverse, and a
    # third application (which by F^4 = I should equal F^-1 again,
    # but here we just care that norm doesn't drift).
    q = Qreg(4, device="cpu")
    q.init_basis(7)
    apply_qft(q)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))
    apply_qft_inverse(q)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ===========================================================================
# Sub-register
# ===========================================================================


def test_qft_subregister_preserves_outer_qubits() -> None:
    # 3-qubit register in |001> (q0=1, q1=0, q2=0). Apply QFT to the
    # sub-register [start=1, n=2] (qubits 1 and 2). The sub-state is
    # |q2 q1> = |00>, so the sub-QFT produces uniform over the four
    # 2-qubit basis states. Qubit 0 stays at 1, so the full register
    # ends up uniform over indices with q0=1: {1, 3, 5, 7}.
    q = Qreg(3, device="cpu")
    q.init_basis(1)
    apply_qft(q, start=1, n=2)
    tol = amp_tol_for(q.dtype)
    for i in range(8):
        if i & 1 == 1:  # q0 (LSB) must be 1
            assert q.amplitude(i) == pytest.approx(0.5 + 0j, abs=tol), (
                f"sub-QFT amp[{i}] = {q.amplitude(i)}, expect 0.5"
            )
        else:
            assert q.amplitude(i) == pytest.approx(0 + 0j, abs=tol), (
                f"sub-QFT amp[{i}] = {q.amplitude(i)}, expect 0 (q0=0)"
            )


def test_qft_subregister_round_trip() -> None:
    # Inverse must work on sub-registers too. Start with a random-ish
    # superposition on the outer qubit, then QFT + inverse on the
    # inner range; the result should equal the starting state.
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    q.apply_h(0)  # outer qubit in superposition
    q.apply_h(2)  # one inner qubit in superposition
    snapshot = q.amplitudes_copy()
    apply_qft(q, start=1, n=2)
    apply_qft_inverse(q, start=1, n=2)
    tol = amp_tol_for(q.dtype)
    for i in range(8):
        assert q.amplitude(i) == pytest.approx(
            complex(snapshot[i].item()), abs=tol
        ), f"sub-QFT round-trip amp[{i}]"


def test_qft_default_n_covers_full_register() -> None:
    # apply_qft(q) with no start/n should cover the whole register,
    # so QFT|0...0> should give uniform 1/sqrt(2^n).
    q = Qreg(3, device="cpu")
    q.init_basis(0)
    apply_qft(q)  # equivalent to apply_qft(q, start=0, n=3)
    expected = 1.0 / math.sqrt(8)
    tol = amp_tol_for(q.dtype)
    for i in range(8):
        assert q.amplitude(i) == pytest.approx(expected + 0j, abs=tol)


# ===========================================================================
# Validation
# ===========================================================================


def test_apply_qft_rejects_negative_start() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_qft: start=-1"):
        apply_qft(q, start=-1)


def test_apply_qft_rejects_zero_n() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_qft: n=0"):
        apply_qft(q, n=0)


def test_apply_qft_rejects_negative_n() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(ValueError, match=r"^qubit: apply_qft: n=-1"):
        apply_qft(q, n=-1)


def test_apply_qft_rejects_range_overflow() -> None:
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_qft: start=1 \+ n=3 > q.n_qubits=3"
    ):
        apply_qft(q, start=1, n=3)


def test_apply_qft_inverse_rejects_range_overflow() -> None:
    # Same validation surface as the forward direction; the inverse
    # uses the same helper.
    q = Qreg(3, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_qft: start=2 \+ n=2 > q.n_qubits=3"
    ):
        apply_qft_inverse(q, start=2, n=2)


# ===========================================================================
# Method-style parity
# ===========================================================================


def test_method_apply_qft_matches_function() -> None:
    q1 = Qreg(3, device="cpu")
    q2 = Qreg(3, device="cpu")
    q1.init_basis(5)
    q2.init_basis(5)
    apply_qft(q1)
    q2.apply_qft()
    tol = amp_tol_for(q1.dtype)
    for i in range(8):
        assert q1.amplitude(i) == pytest.approx(q2.amplitude(i), abs=tol)


def test_method_apply_qft_inverse_matches_function() -> None:
    q1 = Qreg(3, device="cpu")
    q2 = Qreg(3, device="cpu")
    q1.init_basis(5)
    q2.init_basis(5)
    apply_qft_inverse(q1)
    q2.apply_qft_inverse()
    tol = amp_tol_for(q1.dtype)
    for i in range(8):
        assert q1.amplitude(i) == pytest.approx(q2.amplitude(i), abs=tol)


def test_method_qft_subregister_args() -> None:
    # Confirms that start= and n= flow through the method wrapper.
    q1 = Qreg(3, device="cpu")
    q2 = Qreg(3, device="cpu")
    q1.init_basis(1)
    q2.init_basis(1)
    apply_qft(q1, start=1, n=2)
    q2.apply_qft(start=1, n=2)
    tol = amp_tol_for(q1.dtype)
    for i in range(8):
        assert q1.amplitude(i) == pytest.approx(q2.amplitude(i), abs=tol)


# ===========================================================================
# Device-parametrised sanity (small n to keep MPS path cheap)
# ===========================================================================


def test_qft_uniform_across_devices(device: str) -> None:
    n = 3
    q = Qreg(n, device=device)
    q.init_basis(0)
    apply_qft(q)
    expected = 1.0 / math.sqrt(1 << n)
    tol = amp_tol_for(q.dtype)
    for i in range(1 << n):
        assert q.amplitude(i) == pytest.approx(expected + 0j, abs=tol)


def test_qft_round_trip_across_devices(device: str) -> None:
    n = 2
    q = Qreg(n, device=device)
    q.init_basis(1)
    apply_qft(q)
    apply_qft_inverse(q)
    tol = amp_tol_for(q.dtype)
    assert q.amplitude(1) == pytest.approx(1 + 0j, abs=tol)
    for i in (0, 2, 3):
        assert q.amplitude(i) == pytest.approx(0 + 0j, abs=tol)


================================================================================
FILE: implementation/python/tests/test_grover.py
================================================================================
"""Tests for Grover's amplitude amplification.

The highest-signal cases:

* ``test_grover_one_marked_in_16`` -- default iterations on 4 qubits
  with 1 marked state recovers P > 0.9 (analytic ~0.961 at 3 iters).
* ``test_grover_four_marked_in_16_one_iteration`` -- 4 marked out of
  16 hits total marked-probability ~1.0 after exactly 1 iteration.
* ``test_grover_over_iteration_drops_probability`` -- amplitude
  amplification reverses past the optimum (locks "more iterations
  is not always better" -- a real bug-detector).
"""

from __future__ import annotations

import math

import pytest

from qubit import (
    Qreg,
    amp_tol_for,
    apply_grover,
    prob_tol_for,
)

# ===========================================================================
# Canonical Grover identities
# ===========================================================================


def test_grover_one_marked_in_16() -> None:
    # 4 qubits, 1 marked state, default iterations.
    # Default count = floor(pi/4 * sqrt(16)) = floor(pi) = 3.
    # Analytic P after 3 iters = sin(7 * arcsin(1/4))^2 ~= 0.961.
    n = 4
    target = 10

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_grover(q, n, oracle)
    assert q.prob_of(target) > 0.9, (
        f"1-marked-in-16 default iterations: P[{target}] = "
        f"{q.prob_of(target)} (analytic ~0.961)"
    )


def test_grover_four_marked_in_16_one_iteration() -> None:
    # 4 marked out of 16: theta = arcsin(sqrt(4/16)) = pi/6.
    # After 1 iter: sin(3 * pi/6)^2 = sin(pi/2)^2 = 1.
    # So total marked prob is exactly 1.0 (within tolerance).
    n = 4
    marked = [1, 5, 10, 14]

    def oracle(q: Qreg, _user: object) -> None:
        for m in marked:
            q._amp[m] = -q._amp[m]

    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_grover(q, n, oracle, iterations=1)
    total = sum(q.prob_of(m) for m in marked)
    assert total == pytest.approx(1.0, abs=prob_tol_for(q.dtype) * 4)


def test_grover_two_marked_in_8_three_iterations() -> None:
    # 2 marked out of 8: theta = arcsin(sqrt(2/8)) = arcsin(0.5) = pi/6.
    # After 1 iter: sin(3 * pi/6)^2 = 1. So 1 iteration is optimal here
    # too (not 2; this is a known small-system result). Verifies the
    # iteration count IS adjustable -- the function doesn't always do
    # the same number of rounds.
    n = 3
    marked = [3, 5]

    def oracle(q: Qreg, _user: object) -> None:
        for m in marked:
            q._amp[m] = -q._amp[m]

    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_grover(q, n, oracle, iterations=1)
    total = sum(q.prob_of(m) for m in marked)
    assert total == pytest.approx(1.0, abs=prob_tol_for(q.dtype) * 4)


def test_grover_zero_iterations_keeps_uniform_state() -> None:
    # iterations=0 means: prepare uniform superposition and STOP
    # (no oracle calls, no diffusion). Distribution is the uniform
    # 1/sqrt(N) per amplitude.
    n = 3

    def never_called_oracle(q: Qreg, _user: object) -> None:
        raise AssertionError(
            "iterations=0 must not call the oracle"
        )

    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_grover(q, n, never_called_oracle, iterations=0)
    expected = 1.0 / math.sqrt(1 << n)
    tol = amp_tol_for(q.dtype)
    for i in range(1 << n):
        assert q.amplitude(i) == pytest.approx(expected + 0j, abs=tol)


def test_grover_over_iteration_drops_probability() -> None:
    # Past the optimum, amplitude amplification REVERSES (the rotation
    # angle continues past pi/2). 20 iterations on 1-marked-in-16
    # should land far below the ~0.961 optimum at 3 iterations.
    n = 4
    target = 10

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q_opt = Qreg(n, device="cpu")
    q_opt.init_basis(0)
    apply_grover(q_opt, n, oracle)  # default = 3 iterations

    q_over = Qreg(n, device="cpu")
    q_over.init_basis(0)
    apply_grover(q_over, n, oracle, iterations=20)

    assert q_opt.prob_of(target) > q_over.prob_of(target) + 0.1, (
        f"optimum {q_opt.prob_of(target)} should clearly exceed "
        f"over-iterated {q_over.prob_of(target)}"
    )


def test_grover_preserves_norm() -> None:
    n = 4
    target = 10

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q = Qreg(n, device="cpu")
    q.init_basis(0)
    apply_grover(q, n, oracle)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ===========================================================================
# User payload threads through
# ===========================================================================


def test_grover_user_payload_reaches_oracle() -> None:
    # The oracle should be invoked once per iteration with the same
    # user payload. Use a list-recording oracle to verify.
    received: list[object] = []

    def oracle(q: Qreg, user: object) -> None:
        received.append(user)

    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_grover(q, 2, oracle, user="hello", iterations=3)
    assert received == ["hello", "hello", "hello"]


def test_grover_default_user_is_none() -> None:
    # If the caller omits `user`, the oracle gets None. Allows callers
    # who close over their predicate to ignore the second argument.
    seen: list[object] = []

    def oracle(q: Qreg, user: object) -> None:
        seen.append(user)

    q = Qreg(2, device="cpu")
    q.init_basis(0)
    apply_grover(q, 2, oracle, iterations=2)
    assert seen == [None, None]


def test_grover_user_payload_drives_oracle_behaviour() -> None:
    # Demonstrates the typical use of `user`: pass the marked-state
    # index/set through, so the oracle is a pure function of (q, user).
    n = 4

    def oracle(q: Qreg, user: object) -> None:
        # The user payload is an int in this test; cast explicitly so
        # mypy is happy without relaxing the signature.
        assert isinstance(user, int)
        q._amp[user] = -q._amp[user]

    # Same oracle, different user payloads => different marked items.
    q1 = Qreg(n, device="cpu")
    q1.init_basis(0)
    apply_grover(q1, n, oracle, user=3)
    assert q1.prob_of(3) > 0.9

    q2 = Qreg(n, device="cpu")
    q2.init_basis(0)
    apply_grover(q2, n, oracle, user=11)
    assert q2.prob_of(11) > 0.9


# ===========================================================================
# n_qubits=1 edge case (degenerate but supported)
# ===========================================================================


def test_grover_n_qubits_one_does_not_amplify_but_runs() -> None:
    # For n=1 with 1 marked state out of 2, theta = arcsin(1/sqrt(2)) =
    # pi/4. After one iteration: sin(3 * pi/4)^2 = (sqrt(2)/2)^2 = 0.5.
    # So P(marked) == 0.5 regardless of iteration count -- amplitude
    # amplification doesn't kick in for the M = N/2 case.
    # The function should still EXECUTE cleanly.
    def oracle(q: Qreg, _user: object) -> None:
        q._amp[0] = -q._amp[0]

    q = Qreg(1, device="cpu")
    q.init_basis(0)
    apply_grover(q, 1, oracle, iterations=5)
    assert q.prob_of(0) == pytest.approx(0.5, abs=prob_tol_for(q.dtype) * 4)


# ===========================================================================
# Sub-register: searched qubits are [0, n_qubits) of a larger register
# ===========================================================================


def test_grover_on_subregister_leaves_outer_qubits_uniform() -> None:
    # 4-qubit register, but Grover only operates on the lowest 3 qubits.
    # Qubit 3 (highest) is untouched -- it's in |0> initially and stays
    # there. The search succeeds within the 3-qubit subspace.
    full_n = 4
    search_n = 3
    target = 5  # in [0, 8), so within the 3-qubit search space

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q = Qreg(full_n, device="cpu")
    q.init_basis(0)
    apply_grover(q, search_n, oracle, iterations=2)
    # After Grover, qubit 3 should still be 0 -- so all probability
    # mass is in indices < 8.
    p_low = sum(q.prob_of(i) for i in range(8))
    assert p_low == pytest.approx(1.0, abs=prob_tol_for(q.dtype) * 4)
    # And amp[target] should be amplified compared to its uniform start.
    assert q.prob_of(target) > 1.0 / 8


# ===========================================================================
# Validation
# ===========================================================================


def test_grover_rejects_zero_n_qubits() -> None:
    q = Qreg(2, device="cpu")

    def oracle(_q: Qreg, _user: object) -> None:
        pass

    with pytest.raises(
        ValueError, match=r"^qubit: apply_grover: n_qubits=0"
    ):
        apply_grover(q, 0, oracle)


def test_grover_rejects_negative_n_qubits() -> None:
    q = Qreg(2, device="cpu")

    def oracle(_q: Qreg, _user: object) -> None:
        pass

    with pytest.raises(
        ValueError, match=r"^qubit: apply_grover: n_qubits=-1"
    ):
        apply_grover(q, -1, oracle)


def test_grover_rejects_too_large_n_qubits() -> None:
    q = Qreg(2, device="cpu")

    def oracle(_q: Qreg, _user: object) -> None:
        pass

    with pytest.raises(
        ValueError, match=r"^qubit: apply_grover: n_qubits=3"
    ):
        apply_grover(q, 3, oracle)


def test_grover_rejects_negative_iterations() -> None:
    q = Qreg(2, device="cpu")

    def oracle(_q: Qreg, _user: object) -> None:
        pass

    with pytest.raises(
        ValueError, match=r"^qubit: apply_grover: iterations=-1"
    ):
        apply_grover(q, 2, oracle, iterations=-1)


def test_grover_rejects_non_callable_oracle() -> None:
    q = Qreg(2, device="cpu")
    with pytest.raises(
        TypeError, match=r"^qubit: apply_grover: oracle must be callable"
    ):
        apply_grover(q, 2, "not callable")  # type: ignore[arg-type]


# ===========================================================================
# Method-style parity
# ===========================================================================


def test_method_apply_grover_matches_function() -> None:
    n = 3
    target = 5

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q1 = Qreg(n, device="cpu")
    q2 = Qreg(n, device="cpu")
    q1.init_basis(0)
    q2.init_basis(0)
    apply_grover(q1, n, oracle, iterations=2)
    q2.apply_grover(n, oracle, iterations=2)
    tol = amp_tol_for(q1.dtype)
    for i in range(1 << n):
        assert q1.amplitude(i) == pytest.approx(q2.amplitude(i), abs=tol)


def test_method_apply_grover_threads_user_payload() -> None:
    # Verify the method-style API also forwards the user payload.
    received: list[object] = []

    def oracle(q: Qreg, user: object) -> None:
        received.append(user)

    q = Qreg(2, device="cpu")
    q.init_basis(0)
    q.apply_grover(2, oracle, user=42, iterations=2)
    assert received == [42, 42]


# ===========================================================================
# Device-parametrised sanity
# ===========================================================================


def test_grover_one_marked_across_devices(device: str) -> None:
    # 1-marked-in-16 with default iterations on every available device.
    # MPS gets the same amplification math as CPU.
    n = 4
    target = 10

    def oracle(q: Qreg, _user: object) -> None:
        q._amp[target] = -q._amp[target]

    q = Qreg(n, device=device)
    q.init_basis(0)
    apply_grover(q, n, oracle)
    # Slightly looser bound on MPS due to complex64 precision; the
    # 0.9 threshold accommodates either complex64 or complex128.
    assert q.prob_of(target) > 0.9


================================================================================
FILE: implementation/python/tests/test_shor.py
================================================================================
"""Tests for Shor's algorithm: modular exponentiation, period finding, factoring.

The four highest-signal tests:

* ``test_modular_exp_orbit_a2_mod5`` -- locks the (x, y) -> (x, a^x*y mod N)
  permutation against a hand-computed orbit table for the smallest
  non-trivial case.
* ``test_modular_exp_y_ge_N_passes_through`` -- locks the reversibility
  pass-through for y >= N.
* ``test_shor_factor_15_seeded`` -- end-to-end: factor 15 returns
  {3, 5} reproducibly with a fixed seed.
* ``test_shor_period_a2_mod21`` (gated by ``RUN_SHOR_21=1``) -- the
  canonical 16-qubit acceptance test: period of 2 mod 21 must divide 6.
"""

from __future__ import annotations

import os

import pytest

from qubit import (
    Qreg,
    amp_tol_for,
    apply_modular_exp,
    apply_shor_period,
    prob_tol_for,
    shor_factor,
)

# ===========================================================================
# apply_modular_exp: deterministic orbit table on basis states
# ===========================================================================


def test_modular_exp_orbit_a2_mod5() -> None:
    # a=2 mod 5 has period 4: 2^0=1, 2^1=2, 2^2=4, 2^3=3 (then 2^4=1).
    # Layout: counting at [0..1] (t=2), target at [2..4] (n=3), total=5.
    # For y=1 and each x in [0, 4), the modexp gate must move the
    # amplitude from (x, y=1) to (x, y=2^x mod 5).
    expected = [1, 2, 4, 3]
    for x in range(4):
        q = Qreg(5, device="cpu")
        initial = (1 << 2) | (x << 0)  # y=1, x=x
        q.init_basis(initial)
        apply_modular_exp(q, 0, 2, 2, 3, 2, 5)
        final = (expected[x] << 2) | (x << 0)
        assert q.amplitude(final) == pytest.approx(
            1 + 0j, abs=amp_tol_for(q.dtype)
        ), f"a=2 mod 5: (x={x}, y=1) -> (x, y={expected[x]})"


def test_modular_exp_y_ge_N_passes_through() -> None:
    # Layout: 1 counting + 5 target. N=5, so y in [5, 31] passes through.
    # Place amplitude at (x=0, y=10); after modexp, it should remain there.
    q = Qreg(6, device="cpu")
    q.init_basis(10 << 1)  # y=10, x=0
    apply_modular_exp(q, 0, 1, 1, 5, 2, 5)
    assert q.amplitude(10 << 1) == pytest.approx(
        1 + 0j, abs=amp_tol_for(q.dtype)
    )


@pytest.mark.parametrize("y", [5, 6, 7, 15, 31])
def test_modular_exp_y_ge_N_passes_through_each(y: int) -> None:
    # Sweep more y >= N values to ensure the pass-through region is
    # actually identity, not just the one tested point.
    q = Qreg(6, device="cpu")
    q.init_basis(y << 1)  # x=0, target = y
    apply_modular_exp(q, 0, 1, 1, 5, 2, 5)
    assert q.amplitude(y << 1) == pytest.approx(
        1 + 0j, abs=amp_tol_for(q.dtype)
    )


def test_modular_exp_preserves_norm_on_superposition() -> None:
    # Build a uniform superposition over the counting register, leave
    # y=1 in the target, then apply modexp. The result is a uniform
    # superposition over (x, y=2^x mod N) for x in [0, 2^t), so the
    # norm must still be 1.
    n = 4
    t = 4
    q = Qreg(t + n, device="cpu")
    q.init_basis(1 << t)  # x=0, y=1
    for i in range(t):
        q.apply_h(i)
    apply_modular_exp(q, 0, t, t, n, 7, 15)
    assert q.norm() == pytest.approx(1.0, abs=prob_tol_for(q.dtype))


# ===========================================================================
# apply_modular_exp: validation
# ===========================================================================


def test_modular_exp_rejects_t_zero() -> None:
    q = Qreg(5, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_modular_exp: t=0"
    ):
        apply_modular_exp(q, 0, 0, 2, 3, 2, 5)


def test_modular_exp_rejects_n_zero() -> None:
    q = Qreg(5, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_modular_exp: n=0"
    ):
        apply_modular_exp(q, 0, 2, 2, 0, 2, 5)


def test_modular_exp_rejects_counting_overflow() -> None:
    q = Qreg(5, device="cpu")
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_modular_exp: counting range \[3, 7\) out of \[0, 5\)",
    ):
        apply_modular_exp(q, 3, 4, 0, 1, 2, 5)


def test_modular_exp_rejects_target_overflow() -> None:
    q = Qreg(5, device="cpu")
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_modular_exp: target range \[3, 7\) out of \[0, 5\)",
    ):
        apply_modular_exp(q, 0, 1, 3, 4, 2, 5)


def test_modular_exp_rejects_overlap() -> None:
    q = Qreg(8, device="cpu")
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_modular_exp: counting \[0, 4\) and target \[2, 6\) overlap",
    ):
        apply_modular_exp(q, 0, 4, 2, 4, 2, 5)


def test_modular_exp_rejects_small_N() -> None:
    q = Qreg(5, device="cpu")
    with pytest.raises(
        ValueError, match=r"^qubit: apply_modular_exp: N=1 must be >= 2"
    ):
        apply_modular_exp(q, 0, 2, 2, 3, 2, 1)


def test_modular_exp_rejects_oversized_N() -> None:
    q = Qreg(5, device="cpu")
    # n=3 target -> N must be <= 8. N=9 is over.
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_modular_exp: N=9 exceeds target capacity 2\^3",
    ):
        apply_modular_exp(q, 0, 2, 2, 3, 2, 9)


def test_modular_exp_rejects_non_coprime() -> None:
    q = Qreg(5, device="cpu")
    # gcd(6, 15) = 3 != 1.
    with pytest.raises(
        ValueError,
        match=r"^qubit: apply_modular_exp: gcd\(a=6, N=15\) != 1",
    ):
        apply_modular_exp(q, 0, 1, 1, 4, 6, 15)


# ===========================================================================
# apply_shor_period: stochastic but seeded
# ===========================================================================


def test_shor_period_a7_mod15_seeded() -> None:
    # True period of 7 mod 15 is 4. The recovered period must be a
    # divisor of 4, i.e., in {1, 2, 4}. With a fixed seed the run is
    # bit-for-bit reproducible.
    n = 4
    t = 9
    q = Qreg(t + n, device="cpu", seed=42)
    res = apply_shor_period(q, n, t, 0, n, 7, 15)
    assert res.period in {1, 2, 4}, (
        f"recovered period {res.period} not a divisor of true r=4 "
        f"(raw measured = {res.measured})"
    )


def test_shor_period_a7_mod15_multiple_seeds() -> None:
    # The function is stochastic; sweep a handful of seeds and confirm
    # every run produces either a divisor of the true order (4 for
    # a=7 mod 15) OR the period=0 "outright failure" sentinel that
    # apply_shor_period returns on a c=0 readout. If the implementation
    # were wrong (e.g. swapped LSB/MSB on the counting register,
    # off-by-one in the modexp permutation), at least one seed would
    # surface a value outside this set.
    n = 4
    t = 9
    for seed in (1, 7, 42, 100, 2026):
        q = Qreg(t + n, device="cpu", seed=seed)
        res = apply_shor_period(q, n, t, 0, n, 7, 15)
        assert res.period in {0, 1, 2, 4}, (
            f"seed={seed}: period={res.period} (measured={res.measured})"
        )


# ===========================================================================
# shor_factor: end-to-end
# ===========================================================================


def test_shor_factor_15_seeded() -> None:
    res = shor_factor(15, seed=42)
    # Both factors must multiply to 15 and lie in {3, 5}.
    assert {res.p, res.q} == {3, 5}, (
        f"shor_factor(15, seed=42) -> p={res.p}, q={res.q}; "
        f"attempts={res.attempts}"
    )
    assert res.p * res.q == 15
    assert res.attempts >= 1


def test_shor_factor_15_multiple_seeds() -> None:
    # The algorithm is stochastic; with N=15 and 20 attempts, success
    # is essentially guaranteed. Verify across several seeds.
    for seed in (1, 7, 42, 99, 2026):
        res = shor_factor(15, seed=seed)
        assert {res.p, res.q} == {3, 5}, (
            f"seed={seed}: p={res.p}, q={res.q}, attempts={res.attempts}"
        )


def test_shor_factor_even_N_short_circuits() -> None:
    # Even N is factored without any quantum step (attempts == 0).
    for N, expected_p, expected_q in [
        (4, 2, 2),
        (6, 2, 3),
        (14, 2, 7),
        (100, 2, 50),
    ]:
        res = shor_factor(N)
        assert (res.p, res.q) == (expected_p, expected_q), (
            f"shor_factor({N}): p={res.p}, q={res.q}"
        )
        assert res.attempts == 0


def test_shor_factor_seeded_is_deterministic() -> None:
    # Same seed -> same factor pair, same attempt count.
    res_a = shor_factor(15, seed=42)
    res_b = shor_factor(15, seed=42)
    assert res_a == res_b


def test_shor_factor_rejects_small_N() -> None:
    with pytest.raises(
        ValueError, match=r"^qubit: shor_factor: N=1 must be >= 2"
    ):
        shor_factor(1)


def test_shor_factor_rejects_negative_max_attempts() -> None:
    with pytest.raises(
        ValueError, match=r"^qubit: shor_factor: max_attempts=0"
    ):
        shor_factor(15, max_attempts=0)


# ===========================================================================
# Shor-21: 16-qubit period-finding test, gated by RUN_SHOR_21
# ===========================================================================


@pytest.mark.skipif(
    os.environ.get("RUN_SHOR_21") != "1",
    reason="set RUN_SHOR_21=1 to run the 16-qubit Shor-21 test",
)
def test_shor_period_a2_mod21_gated() -> None:
    # The 2 mod 21 orbit is {1, 2, 4, 8, 16, 11, 22%21=1}, so the true
    # period is 6. Continued-fraction recovery returns the denominator
    # of the best rational approximation of c/2^t with denominator <= 21,
    # which must divide 6 -- i.e., be in {1, 2, 3, 6}. The test uses
    # a FIXED a=2 (rather than running shor_factor) so the only
    # stochasticity is the quantum measurement; gated by RUN_SHOR_21=1
    # to keep the default suite fast.
    n = 5
    t = 11
    q = Qreg(t + n, device="cpu", seed=2026)
    res = apply_shor_period(q, n, t, 0, n, 2, 21)
    assert res.period in {1, 2, 3, 6}, (
        f"recovered period {res.period} not a divisor of true r=6 "
        f"(raw measured = {res.measured})"
    )


# ===========================================================================
# Method-style parity
# ===========================================================================


def test_method_apply_modular_exp_matches_function() -> None:
    # The deterministic modexp orbit case, run through both APIs.
    q1 = Qreg(5, device="cpu")
    q2 = Qreg(5, device="cpu")
    q1.init_basis((1 << 2) | (3 << 0))  # y=1, x=3
    q2.init_basis((1 << 2) | (3 << 0))
    apply_modular_exp(q1, 0, 2, 2, 3, 2, 5)
    q2.apply_modular_exp(0, 2, 2, 3, 2, 5)
    tol = amp_tol_for(q1.dtype)
    for i in range(1 << 5):
        assert q1.amplitude(i) == pytest.approx(q2.amplitude(i), abs=tol)


def test_method_apply_shor_period_matches_function() -> None:
    # Both call paths share the same RNG state initially, so the
    # measurement outcome (and therefore the recovered period) must
    # match exactly.
    n = 4
    t = 9
    q1 = Qreg(t + n, device="cpu", seed=42)
    q2 = Qreg(t + n, device="cpu", seed=42)
    res_function = apply_shor_period(q1, n, t, 0, n, 7, 15)
    res_method = q2.apply_shor_period(n, t, 0, n, 7, 15)
    assert res_function == res_method


# ===========================================================================
# Device-parametrised: modular_exp on small register on every device
# ===========================================================================


def test_modular_exp_orbit_across_devices(device: str) -> None:
    # The small a=2 mod 5 orbit on every available device. Verifies
    # the permutation tensor transfers correctly and gather works on
    # both CPU and MPS (and CUDA if present).
    expected = [1, 2, 4, 3]
    for x in range(4):
        q = Qreg(5, device=device)
        initial = (1 << 2) | (x << 0)
        q.init_basis(initial)
        apply_modular_exp(q, 0, 2, 2, 3, 2, 5)
        final = (expected[x] << 2) | (x << 0)
        assert q.prob_of(final) == pytest.approx(
            1.0, abs=prob_tol_for(q.dtype)
        )


================================================================================
FILE: implementation/python/tests/test_cli.py
================================================================================
"""Tests for the qubit-demo CLI.

Drives :func:`qubit.cli.main` directly (no subprocess) so the tests
stay fast and the captured stdout is straightforward to assert on.
``--help`` and unknown-algo paths both exit through argparse's
``SystemExit``; demo paths return integer status codes.
"""

from __future__ import annotations

import pytest

from qubit.cli import main


def test_cli_help_exits_zero(capsys: pytest.CaptureFixture[str]) -> None:
    # argparse's --help prints to stdout and SystemExits with code 0.
    with pytest.raises(SystemExit) as exc_info:
        main(["--help"])
    assert exc_info.value.code == 0
    captured = capsys.readouterr()
    assert "qubit-demo" in captured.out
    assert "--algo" in captured.out


def test_cli_default_algo_is_bell(capsys: pytest.CaptureFixture[str]) -> None:
    code = main([])
    assert code == 0
    out = capsys.readouterr().out
    assert "Bell" in out
    assert "P(00)" in out


@pytest.mark.parametrize("algo", ["bell", "qft", "grover", "shor"])
def test_cli_each_demo_runs(
    capsys: pytest.CaptureFixture[str], algo: str
) -> None:
    code = main(["--algo", algo])
    assert code == 0
    out = capsys.readouterr().out
    assert out.strip() != "", f"{algo} demo produced no output"


def test_cli_bell_output_shape(capsys: pytest.CaptureFixture[str]) -> None:
    # Pin the specific numbers so any regression in the Bell-state
    # pipeline (H, CNOT, prob_of) gets caught here too.
    code = main(["--algo", "bell"])
    assert code == 0
    out = capsys.readouterr().out
    assert "P(00) = 0.5000" in out
    assert "P(11) = 0.5000" in out


def test_cli_shor_output_shape(capsys: pytest.CaptureFixture[str]) -> None:
    # The Shor demo is seeded -> deterministic; pin the factors.
    code = main(["--algo", "shor"])
    assert code == 0
    out = capsys.readouterr().out
    # The seeded run lands on attempt 1 with p=3, q=5 (order
    # depends on which of a^(r/2) +/- 1 gcd-resolves first).
    assert "p=3" in out or "p=5" in out
    assert "q=5" in out or "q=3" in out


def test_cli_unknown_algo_rejected(
    capsys: pytest.CaptureFixture[str],
) -> None:
    # argparse rejects unknown choices and SystemExits with code 2.
    with pytest.raises(SystemExit) as exc_info:
        main(["--algo", "nonsense"])
    assert exc_info.value.code == 2
    captured = capsys.readouterr()
    assert "invalid choice" in captured.err


def test_cli_handles_unexpected_exception(
    capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
) -> None:
    # If a demo function raises, the CLI catches at the top level,
    # prints a qubit-demo-prefixed error to stderr, and returns 1.
    def boom() -> None:
        raise RuntimeError("qubit: simulated failure")

    monkeypatch.setitem(
        __import__("qubit.cli", fromlist=["_DEMOS"])._DEMOS,
        "bell",
        boom,
    )
    code = main(["--algo", "bell"])
    assert code == 1
    err = capsys.readouterr().err
    assert "qubit-demo:" in err
    assert "simulated failure" in err