"""Finite-dimensional recovery illustration using only the Python standard library.

Order of registers is R, E, L, M, with the leftmost bit most significant.
R is an untouched reference, E and L are collected radiation registers and
M is an inaccessible record. The physical interpretation of those labels
is an assumption. This script specifies no gravitational dynamics.
"""

import json
import math


def reduced_density(state, retained):
    """Trace the four-qubit pure state over registers outside retained."""
    discarded = [register for register in range(4) if register not in retained]
    dimension = 2 ** len(retained)
    density = [[0j for _ in range(dimension)] for _ in range(dimension)]

    def label(index, registers):
        value = 0
        for register in registers:
            value = 2 * value + ((index >> (3 - register)) & 1)
        return value

    for row, amplitude in enumerate(state):
        for column, other in enumerate(state):
            if label(row, discarded) == label(column, discarded):
                density[label(row, retained)][label(column, retained)] += (
                    amplitude * other.conjugate()
                )
    return density


def encoded_state(alpha, beta, overlap, reference=False):
    """Map |i> to |i>_E |i>_L |m_i>_M, retaining R when requested.

    |m_0> = |0> and |m_1> = overlap|0> + sqrt(1-overlap**2)|1>.
    Orthogonality of the E,L codewords makes this an isometry.
    """
    state = [0j] * 16
    reference_bit = 8 if reference else 0
    state[0] = complex(alpha)
    state[reference_bit + 6] = complex(beta * overlap)
    state[reference_bit + 7] = complex(beta * math.sqrt(1 - overlap**2))
    return state


def decode(state):
    """Apply CNOT E -> L after both radiation registers are collected."""
    decoded = [0j] * 16
    for index, amplitude in enumerate(state):
        destination = index ^ 2 if index & 4 else index
        decoded[destination] = amplitude
    return decoded


def close(actual, expected):
    assert abs(actual - expected) < 1e-12, (actual, expected)


def calculate(overlap):
    amplitude = 1 / math.sqrt(2)
    positive = encoded_state(amplitude, amplitude, overlap)
    negative = encoded_state(amplitude, -amplitude, overlap)
    radiation_positive = reduced_density(positive, [1, 2])
    radiation_negative = reduced_density(negative, [1, 2])

    # The difference has only two nonzero off-diagonal entries. Its
    # eigenvalues are +|entry| and -|entry|, so D = |entry|.
    phase_distance = abs(radiation_positive[0][3] - radiation_negative[0][3])
    close(phase_distance, overlap)
    for retained in ([1], [2]):
        first = reduced_density(positive, retained)
        second = reduced_density(negative, retained)
        for row in range(2):
            for column in range(2):
                close(first[row][column], second[row][column])

    entangled = encoded_state(amplitude, amplitude, overlap, reference=True)
    recovered = reduced_density(decode(entangled), [0, 1])
    fidelity = sum(recovered[row][column] for row in (0, 3) for column in (0, 3)) / 2
    close(fidelity, (1 + overlap) / 2)

    # Check the entire decoded one-qubit channel for different populations
    # and phases, beyond the two phase inputs used in the explanation.
    for theta in (0, 0.37, 1.1, math.pi / 2):
        for phase in (0, 0.71, math.pi):
            alpha = math.cos(theta)
            beta = math.sin(theta) * complex(math.cos(phase), math.sin(phase))
            output = reduced_density(decode(encoded_state(alpha, beta, overlap)), [1])
            expected = [
                [abs(alpha)**2, overlap * alpha * beta.conjugate()],
                [overlap * beta * alpha, abs(beta)**2],
            ]
            for row in range(2):
                for column in range(2):
                    close(output[row][column], expected[row][column])

    return {
        "remnant_state_overlap": overlap,
        "joint_radiation_phase_trace_distance": phase_distance,
        "decoded_reference_entanglement_fidelity": fidelity.real,
        "remnant_branch_trace_distance": math.sqrt(1 - overlap**2),
    }


if __name__ == "__main__":
    grid = [calculate(index / 10) for index in range(11)]
    print(json.dumps({
        "scope": "Ideal finite-dimensional isometry with a specified decoder",
        "overlap_domain": "Real g in [0,1]",
        "decoder": "CNOT E -> L, recover logical qubit on E",
        "checks": "Reduced states, reference fidelity and decoded channel on a population/phase grid",
        "physical_model": "Gravitational geometry, energy budget and emission dynamics require specification",
        "results": grid,
    }, indent=2))