"""Independent numerical checks for the proposed relational-time benchmark.

These are theoretical consistency checks, not experimental results. No site edits.
Units: hbar = omega = 1. Standard-library + numpy only.
"""
import json
from pathlib import Path
import numpy as np


def phase_state(n, theta):
    return np.exp(-1j * np.arange(n) * theta) / np.sqrt(n)


def history(n):
    result = np.zeros((n, n), dtype=complex)
    for m in range(n):
        result[n - 1 - m, m] = 1 / np.sqrt(n)
    return result


def condition(n, theta):
    result = phase_state(n, theta).conj() @ history(n)
    return result / np.linalg.norm(result)


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


grid = np.linspace(0, 2*np.pi, 361)
constraint_errors = []
conditional_errors = []
povm_errors = []
for n in (2, 3, 4, 5, 8):
    h = np.diag(np.arange(n))
    c = np.kron(h, np.eye(n)) + np.kron(np.eye(n), h) - (n-1)*np.eye(n*n)
    constraint_errors.append(float(np.linalg.norm(c @ history(n).reshape(-1))))
    for theta in grid:
        expected = phase_state(n, theta)
        conditional_errors.append(float(np.linalg.norm(density(condition(n, theta))-density(expected))))
    # Discrete Fourier quadrature exactly integrates the phase POVM for this n.
    ph = np.arange(4*n) * (2*np.pi/(4*n))
    integrated = sum(n*density(phase_state(n,t))/(4*n) for t in ph)
    povm_errors.append(float(np.linalg.norm(integrated-np.eye(n))))

V = 0.8
sigma = float(np.sqrt(-2*np.log(V)))
a = float(np.arccos(V))
phi = np.linspace(-np.pi, np.pi, 8192, endpoint=False)
# Independently perform wrapped-Gaussian quadrature; enough images at this sigma.
w = sum(np.exp(-((phi+2*np.pi*k)/sigma)**2/2) for k in range(-4,5))
w /= w.sum()
qutrit_errors = []
qubit_difference = []
min_eigenvalue = 1.0
observed = {}
for theta in grid:
    states = np.exp(-1j*(theta+phi[:,None])*np.arange(3)[None,:])/np.sqrt(3)
    rho_g = np.einsum('i,ij,ik->jk', w, states, states.conj())
    rho_b = (density(phase_state(3,theta+a))+density(phase_state(3,theta-a)))/2
    initial = phase_state(3,0)
    p_g = float(np.real(initial.conj()@rho_g@initial))
    p_b = float(np.real(initial.conj()@rho_b@initial))
    expected_g = (3+4*V*np.cos(theta)+2*V**4*np.cos(2*theta))/9
    expected_b = (3+4*V*np.cos(theta)+2*(2*V**2-1)*np.cos(2*theta))/9
    qutrit_errors.extend([abs(p_g-expected_g), abs(p_b-expected_b)])
    min_eigenvalue = min(min_eigenvalue, float(np.linalg.eigvalsh(rho_g).min()),float(np.linalg.eigvalsh(rho_b).min()))
    p2_g = float(np.sum(w*(1+np.cos(theta+phi))/2))
    p2_b = ((1+np.cos(theta+a))/2+(1+np.cos(theta-a))/2)/2
    qubit_difference.append(abs(p2_g-p2_b))
    if theta == 0:
        observed = {'wrapped_gaussian_return':p_g, 'binary_offset_return':p_b, 'difference':p_g-p_b}

lam = .02
h = np.diag(np.arange(3))
interaction_constraint = np.kron(h,np.eye(3))+np.kron(np.eye(3),h)-lam*np.kron(h,h)-2*np.eye(9)
interaction_residual = float(np.linalg.norm(interaction_constraint@history(3).reshape(-1)))
results = {
    'status':'theoretical numerical checks only; no experimental data',
    'max_constraint_residual':max(constraint_errors),
    'max_conditional_density_error':max(conditional_errors),
    'max_phase_povm_normalization_error':max(povm_errors),
    'max_qutrit_formula_vs_independent_quadrature_error':max(qutrit_errors),
    'max_qubit_difference_between_matched_visibility_models':max(qubit_difference),
    'minimum_density_eigenvalue_with_roundoff':min_eigenvalue,
    'noise_parameters':{'V':V,'wrapped_gaussian_sigma_radians':sigma,'binary_offset_radians':a},
    'qutrit_at_phase_zero':observed,
    'unchanged_history_with_interaction':{'lambda':lam,'residual':interaction_residual,'expected_residual':lam/np.sqrt(3)},
}
assert max(constraint_errors) < 1e-12
assert max(conditional_errors) < 1e-12
assert max(povm_errors) < 1e-12
assert max(qutrit_errors) < 1e-12
assert max(qubit_difference) < 1e-12
assert min_eigenvalue > -1e-12
assert abs(observed['difference']-.0288) < 1e-12
assert abs(interaction_residual-lam/np.sqrt(3)) < 1e-12
destination = Path(__file__).with_name('timeless-model-checks.json')
destination.write_text(json.dumps(results,indent=2)+'\n')
print(json.dumps(results,indent=2))