"""Reproducible local examples for Introduction to Quantum Computing — 2026.
Tested with Python 3.12 and Qiskit 2.5.2. No cloud access is used.
"""
import json
import numpy as np
import qiskit
from qiskit import QuantumCircuit
from qiskit.primitives import StatevectorSampler, StatevectorEstimator
from qiskit.quantum_info import Statevector, Pauli, DensityMatrix, partial_trace, SparsePauliOp, Kraus

def main():
    psi = Statevector(np.array([2, 3j]) / np.sqrt(13))
    expectations = {axis: float(psi.expectation_value(Pauli(axis)).real) for axis in ("X", "Y", "Z")}
    np.testing.assert_allclose(list(expectations.values()), [0, 12/13, -5/13], atol=1e-12)
    for basis, expected in [("Z", [4/13, 9/13]), ("X", [.5, .5]), ("Y", [25/26, 1/26])]:
        qc = QuantumCircuit(1)
        if basis == "Y":
            qc.sdg(0)
        if basis != "Z":
            qc.h(0)
        np.testing.assert_allclose(psi.evolve(qc).probabilities(), expected, atol=1e-12)
    bell = QuantumCircuit(2)
    bell.h(0)
    bell.cx(0, 1)
    state = Statevector.from_instruction(bell)
    reduced = partial_trace(DensityMatrix(state), [1])
    np.testing.assert_allclose(reduced.data, np.eye(2)/2, atol=1e-12)
    asymmetric = QuantumCircuit(3)
    asymmetric.x(2)  # book q0 -> SDK wire n-1-0
    assert np.argmax(Statevector.from_instruction(asymmetric).probabilities()) == 4
    measured = bell.copy()
    measured.measure_all()
    counts = StatevectorSampler(seed=2026).run([measured], shots=2048).result()[0].data.meas.get_counts()
    assert set(counts).issubset({"00", "11"}) and sum(counts.values()) == 2048
    observable = SparsePauliOp.from_list([("XX", 1.), ("ZZ", 1.)])
    ev = StatevectorEstimator().run([(bell, observable)]).result()[0].data.evs
    np.testing.assert_allclose(ev, 2, atol=1e-12)
    gamma = .25
    ks = [np.diag([1, np.sqrt(1-gamma)]), np.array([[0, np.sqrt(gamma)], [0, 0]])]
    damped = DensityMatrix(Statevector([1, 1])/np.sqrt(2)).evolve(Kraus(ks))
    np.testing.assert_allclose(damped.data, [[5/8, np.sqrt(3)/4], [np.sqrt(3)/4, 3/8]], atol=1e-12)
    print(json.dumps({"qiskit": qiskit.__version__, "expectations": expectations,
                      "bell_counts": counts, "bell_XX_plus_ZZ": float(ev),
                      "checks": "passed"}, indent=2))

if __name__ == "__main__":
    main()