import cmath
import math

def normalize(state):
    norm = math.sqrt(sum(abs(value) ** 2 for value in state))
    if norm < 1e-15:
        raise ValueError("the zero vector is not a state")
    return tuple(value / norm for value in state)

def state_to_bloch(state):
    alpha, beta = normalize(state)
    coherence = alpha.conjugate() * beta
    return 2 * coherence.real, 2 * coherence.imag, abs(alpha) ** 2 - abs(beta) ** 2

def bloch_to_state(point):
    x, y, z = point
    radius = math.sqrt(x * x + y * y + z * z)
    if abs(radius - 1.0) > 1e-12:
        raise ValueError("pure-state Bloch vectors must have unit radius")
    if z < -1 + 1e-12:
        return 0j, 1 + 0j
    alpha = math.sqrt((1 + z) / 2)
    return alpha, complex(x, y) / (2 * alpha)

def matvec(matrix, state):
    return tuple(sum(matrix[row][column] * state[column] for column in range(2)) for row in range(2))

def fidelity(left, right):
    left, right = normalize(left), normalize(right)
    overlap = sum(a.conjugate() * b for a, b in zip(left, right))
    return abs(overlap) ** 2

scale = math.sqrt(0.5)
states = []
for theta, phi in ((0, 0), (math.pi, 0), (math.pi / 3, 0.7), (math.pi / 2, -1.2), (2.2, 2.6)):
    states.append((math.cos(theta / 2), cmath.exp(1j * phi) * math.sin(theta / 2)))
states.append(tuple(cmath.exp(0.91j) * value for value in states[2]))

max_roundtrip_infidelity = 0.0
for state in states:
    point = state_to_bloch(state)
    assert abs(sum(axis * axis for axis in point) - 1.0) < 1e-12
    recovered = bloch_to_state(point)
    infidelity = 1.0 - fidelity(state, recovered)
    max_roundtrip_infidelity = max(max_roundtrip_infidelity, abs(infidelity))
assert max_roundtrip_infidelity < 1e-12
assert all(abs(a - b) < 1e-12 for a, b in zip(state_to_bloch(states[2]), state_to_bloch(states[-1])))

I = ((1, 0), (0, 1))
X = ((0, 1), (1, 0))
Y = ((0, -1j), (1j, 0))
Z = ((1, 0), (0, -1))
H = ((scale, scale), (scale, -scale))
S = ((1, 0), (0, 1j))
rotations = {
    "I": (I, lambda x, y, z: (x, y, z)),
    "X": (X, lambda x, y, z: (x, -y, -z)),
    "Y": (Y, lambda x, y, z: (-x, y, -z)),
    "Z": (Z, lambda x, y, z: (-x, -y, z)),
    "H": (H, lambda x, y, z: (z, -y, x)),
    "S": (S, lambda x, y, z: (-y, x, z)),
}
max_rotation_error = 0.0
for matrix, rotation in rotations.values():
    for state in states:
        actual = state_to_bloch(matvec(matrix, state))
        expected = rotation(*state_to_bloch(state))
        max_rotation_error = max(max_rotation_error, *(abs(a - b) for a, b in zip(actual, expected)))
assert max_rotation_error < 1e-12

mixed_point_rejected = False
try:
    bloch_to_state((0.2, 0.0, 0.0))
except ValueError:
    mixed_point_rejected = True
assert mixed_point_rejected
wrong_y = lambda state: -2 * (state[0].conjugate() * state[1]).imag
assert abs(wrong_y(states[3]) - state_to_bloch(states[3])[1]) > 0.5
print(f"PASS: 17 Bloch viewer round-trips {len(states)} states through {len(rotations)} gate rotations; last |infidelity|={abs(infidelity):.2e}, max infidelity={max_roundtrip_infidelity:.2e}, max SO(3) error={max_rotation_error:.2e}")
