"""Reproduce the ideal QMM phase-recovery comparison. Run with Python 3 and NumPy.

The register order is field Q0, memory Q1 and output Q2. This script
constructs Eq. (8), evaluates the proposed CNOT decoder and traces out
Q0 and Q1. All gates are ideal.
"""

import numpy as np

phi = np.pi / 6
zero, one = np.eye(2, dtype=complex)
phase_rotation = np.diag(np.exp(np.array([-1j, 1j]) * phi / 2))
identity = np.eye(2)
controlled_swap = np.eye(8)[[0, 1, 2, 3, 4, 6, 5, 7]]
decode = np.eye(8)[[0, 5, 2, 7, 4, 1, 6, 3]]

def make_cycle(theta):
    c, s = np.cos(theta / 2), np.sin(theta / 2)
    controlled_ry = np.eye(8)
    controlled_ry[4:, 4:] = np.kron([[c, -s], [s, c]], identity)
    return (
        np.kron(np.eye(4), phase_rotation.conj().T)
        @ controlled_swap
        @ np.kron(np.kron(identity, phase_rotation), identity)
        @ controlled_ry
    )

def output(cycle, state):
    joint = cycle @ np.kron(np.kron(state, zero), zero)
    amplitudes = joint.reshape(4, 2)
    return amplitudes.T @ amplitudes.conj()

def projector(state):
    return np.outer(state, state.conj())

def trace_distance(left, right):
    return 0.5 * np.abs(np.linalg.eigvalsh(left - right)).sum()

states = [zero, one] + [(zero + relative_phase * one) / np.sqrt(2) for relative_phase in (1, -1, 1j, -1j)]
theta = np.pi / 4
cycle = make_cycle(theta)
assert np.allclose(cycle.conj().T @ cycle, np.eye(8))
channels = {
    "QMM Eq. 8": lambda state: output(cycle, state),
    "QMM with CNOT at pi/4": lambda state: output(decode @ cycle, state),
    "QMM with CNOT at pi/2": lambda state: output(decode @ make_cycle(np.pi / 2), state),
    "Z measurement and re-preparation": lambda state: np.diag(abs(state) ** 2),
    "Fixed zero": lambda state: projector(zero),
    "Ideal identity": projector,
}
c, s = np.cos(theta / 2), np.sin(theta / 2)
v = c * zero + s * one
for state in states:
    p_one = abs(state[1]) ** 2
    assert np.allclose(output(cycle, state), (1 - p_one) * projector(zero) + p_one * projector(v))
for label, channel in channels.items():
    outputs = [channel(state) for state in states]
    fidelities = [np.vdot(state, rho @ state).real for state, rho in zip(states, outputs)]
    dx = trace_distance(outputs[2], outputs[3])
    dy = trace_distance(outputs[4], outputs[5])
    print(f"{label:34} {dx:.8f} {dy:.8f} {np.mean(fidelities):.8f}")
for theta in np.linspace(0, np.pi, 101):
    s = np.sin(theta / 2)
    candidate = decode @ make_cycle(theta)
    fidelities = []
    for state in states:
        alpha, beta = state
        p_one = abs(beta) ** 2
        expected = np.array([[1 - p_one * s**2, alpha * beta.conj() * s],
                             [alpha.conj() * beta * s, p_one * s**2]])
        rho = output(candidate, state)
        assert np.allclose(rho, expected)
        fidelities.append(np.vdot(state, rho @ state).real)
    assert np.isclose(np.mean(fidelities), 0.5 + (2 * s + s**2) / 6)
print(f"Ideal threshold: theta > {2 * np.arcsin(np.sqrt(2) - 1):.8f} radians")