import math

def outer(state):
    return tuple(tuple(state[row] * state[column].conjugate() for column in range(len(state)))
                 for row in range(len(state)))

def mixture(weighted_states):
    if abs(sum(weight for weight, _ in weighted_states) - 1.0) > 1e-12:
        raise ValueError("mixture weights must sum to one")
    size = len(weighted_states[0][1])
    return tuple(tuple(sum(weight * state[row] * state[column].conjugate()
                           for weight, state in weighted_states)
                       for column in range(size)) for row in range(size))

def trace(matrix):
    return sum(matrix[index][index] for index in range(len(matrix)))

def purity(matrix):
    return sum(abs(value) ** 2 for row in matrix for value in row).real

def validate_qubit_density(matrix):
    if len(matrix) != 2 or any(len(row) != 2 for row in matrix):
        raise ValueError("fixture expects a qubit density matrix")
    if abs(trace(matrix) - 1.0) > 1e-12 or abs(matrix[1][0] - matrix[0][1].conjugate()) > 1e-12:
        raise ValueError("density matrix must be Hermitian with trace one")
    determinant = (matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0]).real
    if matrix[0][0].real < -1e-12 or matrix[1][1].real < -1e-12 or determinant < -1e-12:
        raise ValueError("density matrix must be positive semidefinite")
    return determinant

def probability(matrix, ket):
    return sum(ket[row].conjugate() * matrix[row][column] * ket[column]
               for row in range(2) for column in range(2)).real

def phase_flip_channel(matrix, probability_z):
    if not 0 <= probability_z <= 1:
        raise ValueError("channel probability must lie in [0,1]")
    factor = 1 - 2 * probability_z
    return ((matrix[0][0], factor * matrix[0][1]),
            (factor * matrix[1][0], matrix[1][1]))

def trace_out_second(matrix):
    if len(matrix) != 4 or any(len(row) != 4 for row in matrix):
        raise ValueError("partial trace fixture expects two qubits")
    return tuple(tuple(sum(matrix[2 * left + bit][2 * right + bit] for bit in (0, 1))
                       for right in (0, 1)) for left in (0, 1))

scale = math.sqrt(0.5)
zero, one, plus, plus_i = (1, 0), (0, 1), (scale, scale), (scale, 1j * scale)
rho_plus = outer(plus)
rho_plus_i = outer(plus_i)
rho_mixed = mixture(((0.5, zero), (0.5, one)))
for matrix in (rho_plus, rho_plus_i, rho_mixed):
    validate_qubit_density(matrix)
assert abs(purity(rho_plus) - 1.0) < 1e-12
assert abs(purity(rho_plus_i) - 1.0) < 1e-12
assert abs(purity(rho_mixed) - 0.5) < 1e-12
assert abs(probability(rho_plus, zero) - probability(rho_mixed, zero)) < 1e-12
x_contrast = probability(rho_plus, plus) - probability(rho_mixed, plus)
assert abs(x_contrast - 0.5) < 1e-12

channel_purities = []
for probability_z in (0.0, 0.25, 0.5):
    evolved = phase_flip_channel(rho_plus, probability_z)
    validate_qubit_density(evolved)
    channel_purities.append(purity(evolved))
assert all(abs(actual - expected) < 1e-12 for actual, expected in zip(channel_purities, (1.0, 0.625, 0.5)))

bell = (scale, 0, 0, scale)
marginal = trace_out_second(outer(bell))
assert all(abs(marginal[row][column] - (0.5 if row == column else 0.0)) < 1e-12
           for row in range(2) for column in range(2))

non_psd_rejected = False
try:
    validate_qubit_density(((0.5, 0.6), (0.6, 0.5)))
except ValueError:
    non_psd_rejected = True
assert non_psd_rejected
bad_weights_rejected = False
try:
    mixture(((0.4, zero), (0.4, one)))
except ValueError:
    bad_weights_rejected = True
assert bad_weights_rejected
print(f"PASS: 15 density laboratory separates coherent/mixed X statistics by {x_contrast:.3f}, tracks dephasing purities={channel_purities}, and recovers Bell marginal purity={purity(marginal):.3f}")