"""Two-qubit information-recovery control, NOT Hawking radiation simulation.

Encode two classical bits as a Bell state. Inspect one emitted subsystem,
then permit access to both qubits and decode coherently. Compare dephasing.
"""
from pathlib import Path
import hashlib
import json
import numpy as np


def main():
    hadamard = np.array([[1,1],[1,-1]], dtype=complex)/np.sqrt(2)
    cnot = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]], dtype=complex)
    unitary = cnot @ np.kron(hadamard, np.eye(2))
    assert np.allclose(unitary.conj().T @ unitary, np.eye(4))
    rows = []
    for message in range(4):
        initial = np.eye(4, dtype=complex)[:,message]
        encoded = unitary @ initial
        amplitudes = encoded.reshape(2,2)
        # Trace out the second qubit: first qubit alone is maximally mixed.
        local = amplitudes @ amplitudes.conj().T
        assert np.allclose(local, np.eye(2)/2)
        recovered = unitary.conj().T @ encoded
        fidelity = float(abs(np.vdot(initial,recovered))**2)
        assert np.isclose(fidelity,1)
        # Unrecorded Z measurement on both qubits destroys phase coherence.
        dephased = np.diag(np.abs(encoded)**2)
        after = unitary.conj().T @ dephased @ unitary
        damaged_success = float(after[message,message].real)
        assert np.isclose(damaged_success,0.5)
        # Canonical probabilities eliminate floating serialization artifacts.
        distribution = [round(float(local[i,i].real),12) for i in range(2)]
        serialized = json.dumps(distribution,separators=(',',':')).encode()
        hash64 = hashlib.shake_256(b'Bell-local-distribution-v1\0'+serialized).hexdigest(8)
        rows.append(dict(message=f'{message:02b}',
            local_measurement_probabilities=distribution,
            global_Z_probabilities=(np.abs(encoded)**2).round(12).tolist(),
            full_coherent_decode_fidelity=fidelity,
            decode_success_after_Z_dephasing=damaged_success,
            local_distribution_hash64=hash64))
    assert len({r['local_distribution_hash64'] for r in rows})==1
    results=dict(scope='Bell encoding and staged subsystem access; no geometry or thermal dynamics',
                 measurement='Exact probabilities, no hardware or random samples',
                 full_decoder_access='Both coherent qubits; no access claimed from one subsystem alone',
                 rows=rows)
    Path(__file__).with_name('entangled-return-results.json').write_text(json.dumps(results,indent=2)+'\n')
    print(json.dumps(results,indent=2))


if __name__=='__main__':
    main()