#!/usr/bin/env python3
"""Generate reproducible Majorana curves and analytic model benchmarks.

Requires Python 3.9+ and NumPy; no network, external data, or SciPy.
Run: python3 generate_majorana_data.py [--output-dir DIRECTORY]
Default output: this script's directory. Writes majorana-data.json and
model-checks.json only after all numerical validation succeeds.

The beta=0 autocorrelations illustrate cancellation under averaging. They are
not an experimental replication, thermalization test, or gravity measurement.
Hamiltonian coefficients are preserved without bandwidth rescaling. Comparing
lambda=0 with lambda=.3 changes both commutation structure and energy scales;
it does not isolate the causal effect of noncommutation.
"""
import argparse
import hashlib
import json
import math
from functools import reduce
from pathlib import Path

import numpy as np

D = 16
TOL = 1e-11
BASE_TERMS = [(-.36, [1, 2, 4, 5]), (.19, [1, 3, 4, 7]),
              (-.71, [1, 3, 5, 6]), (.22, [2, 3, 4, 6]),
              (.49, [2, 3, 5, 7])]
PERTURBATION_INDICES = [1, 2, 3, 5]
I = np.eye(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.diag([1, -1]).astype(complex)


def maxabs(a):
    return float(np.max(np.abs(a)))


def check(name, error, tolerance=TOL):
    error = float(error)
    if not math.isfinite(error) or error > tolerance:
        raise AssertionError(f'{name}: {error} exceeds {tolerance}')
    return {'passed': True, 'max_absolute_error': error, 'tolerance': tolerance}


def majoranas():
    return [reduce(np.kron, [Z] * j + [p] + [I] * (3-j)) / math.sqrt(2)
            for j in range(4) for p in (X, Y)]


def product(psi, indices):
    return reduce(np.matmul, [psi[k-1] for k in indices])


def spectral_traces(H, psi, times):
    E, V = np.linalg.eigh(H)
    gaps = E[:, None] - E[None, :]
    weights = [np.abs(V.conj().T @ p @ V) ** 2 for p in psi]
    curves = np.array([[2 / D * np.sum(w * np.cos(gaps*t)) for t in times]
                       for w in weights], dtype=float)
    return curves, E, V


def matrix_exp_taylor(A):
    """Independent scaling-and-squaring exponential, with converged Taylor sum."""
    norm = float(np.linalg.norm(A, ord=np.inf))
    scale = max(0, math.ceil(math.log2(norm))) if norm else 0
    B = A / (2 ** scale)
    result = np.eye(A.shape[0], dtype=complex)
    term = result.copy()
    for k in range(1, 150):
        term = term @ B / k
        result = result + term
        if np.linalg.norm(term, ord=np.inf) < 1e-17:
            break
    else:
        raise AssertionError('Taylor exponential failed to converge')
    for _ in range(scale):
        result = result @ result
    return result


def majorana_data():
    psi = majoranas()
    identity = np.eye(D)
    terms = [product(psi, indices) for _, indices in BASE_TERMS]
    perturbation = product(psi, PERTURBATION_INDICES)
    H0 = sum(coefficient * term for (coefficient, _), term in zip(BASE_TERMS, terms))
    times = np.arange(241, dtype=float) / 10
    validations = {
        'majorana_anticommutation': check('Majorana algebra', max(
            maxabs(a@b+b@a-(identity if i == j else 0))
            for i, a in enumerate(psi) for j, b in enumerate(psi))),
        'baseline_quartic_terms_commute': check('Baseline term commutation', max(
            maxabs(a@b-b@a) for a in terms for b in terms)),
        'quartic_terms_hermitian': check('Quartic Hermiticity', max(
            maxabs(a-a.conj().T) for a in terms+[perturbation])),
    }
    zero, _, _ = spectral_traces(np.zeros((D, D)), psi, times)
    validations['zero_hamiltonian_all_times'] = check('Zero H', maxabs(zero-1))
    perturbation_commutators = [float(np.linalg.norm(perturbation@a-a@perturbation,
                                                   ord='fro')) for a in terms]
    if not any(v > TOL for v in perturbation_commutators):
        raise AssertionError('Perturbation must fail to commute with some baseline terms')
    cases = []
    for lam in (0., .3):
        H = H0 + lam * perturbation
        curves, E, V = spectral_traces(H, psi, times)
        direct_error = 0.
        unitarity_error = 0.
        for k, t in enumerate(times):
            U = (V * np.exp(-1j*E*t)) @ V.conj().T
            direct = np.array([2/D*np.trace(U.conj().T@p@U@p).real for p in psi])
            direct_error = max(direct_error, maxabs(direct-curves[:, k]))
            unitarity_error = max(unitarity_error, maxabs(U.conj().T@U-identity))
        independent_U_error = 0.
        independent_trace_error = 0.
        for t in (0., .1, 2.8, 10., 20., 24.):
            U = matrix_exp_taylor(-1j*H*t)
            U_spectral = (V * np.exp(-1j*E*t)) @ V.conj().T
            independent_U_error = max(independent_U_error, maxabs(U-U_spectral))
            direct = np.array([2/D*np.trace(U.conj().T@p@U@p).real for p in psi])
            independent_trace_error = max(independent_trace_error,
                                          maxabs(direct-curves[:, round(t*10)]))
        case_checks = {
            'hamiltonian_hermiticity': check('H Hermiticity', maxabs(H-H.conj().T)),
            'spectral_vs_direct_U_all_times': check('Spectral/direct trace', direct_error),
            'unitarity_all_times': check('Unitarity', unitarity_error),
            'independent_taylor_U_selected_times': check('Independent exponential', independent_U_error),
            'independent_taylor_traces_selected_times': check('Independent traces', independent_trace_error),
            'spectator_commutator': check('Spectator commutation', maxabs(H@psi[7]-psi[7]@H)),
            'spectator_all_times': check('Spectator value', maxabs(curves[7]-1)),
            'initial_values': check('t=0', maxabs(curves[:, 0]-1)),
            'correlation_magnitude_bound': check('Magnitude bound', max(0., float(np.max(np.abs(curves)))-1)),
        }
        active = curves[:7]
        case = {
            'perturbation_lambda': lam,
            'traces': np.round(curves, 14).tolist(),
            'mean_active7': np.round(np.mean(active, axis=0), 14).tolist(),
            'mean_absolute_active7': np.round(np.mean(np.abs(active), axis=0), 14).tolist(),
            'mean_all8': np.round(np.mean(curves, axis=0), 14).tolist(),
            'mean_absolute_all8': np.round(np.mean(np.abs(curves), axis=0), 14).tolist(),
            'rms_energy': float(np.sqrt(np.trace(H@H).real/D)),
            'spectral_width': float(E[-1]-E[0]),
            'validation': case_checks,
        }
        cases.append(case)
    benchmark = {
        2.8: [.460738, .645686, .430291, .805103, .369508, .519679, .746572, 1.],
        10.: [.121290, .079378, .002897, -.059947, -.160605, -.416290, -.448030, 1.],
        20.: [.198459, .098431, -.523642, -.170614, -.114495, -.402856, -.060298, 1.],
    }
    baseline = np.array(cases[0]['traces'])
    error = max(maxabs(baseline[:, round(t*10)]-values) for t, values in benchmark.items())
    validations['review_six_decimal_benchmarks'] = check('Review table', error, 5.1e-7)
    snapshots = [{
        'time': t,
        'traces': baseline[:, round(t*10)].tolist(),
        **{key: cases[0][key][round(t*10)] for key in
           ('mean_active7', 'mean_absolute_active7', 'mean_all8', 'mean_absolute_all8')}
    } for t in (0., 2.8, 10., 20.)]
    return {
        'schema_version': 1,
        'metadata': {
            'title': 'What an average hides: exact infinite-temperature autocorrelations',
            'status': 'Reproducible numerical benchmark; no experiment replication claim',
            'state': 'beta = 0; rho = I16 / 16',
            'beta': 0,
            'dimension': D,
            'qubits': 4,
            'majoranas': 8,
            'normalization': 'psi_i = gamma_i / sqrt(2); {psi_i,psi_j} = delta_ij I; hbar = 1',
            'majorana_spin_mapping': 'gamma_(2j-1) = Z tensor ... tensor Z tensor X_j tensor I ...; gamma_(2j) replaces X_j by Y_j; j=1..4, qubit 1 is leftmost tensor factor',
            'observable': 'G_i(t) = (2/16) Re Tr[exp(+iHt) psi_i exp(-iHt) psi_i]',
            'time_units': 'Inverse Hamiltonian coefficient units, hbar=1; no physical seconds assigned',
            'hamiltonian_terms': [{'coefficient': c, 'indices': idx} for c, idx in BASE_TERMS],
            'perturbation': {'coefficient': 'lambda', 'indices': PERTURBATION_INDICES},
            'product_order': 'Listed ascending Majorana indices multiplied left to right',
            'rescaling': 'None. Original coefficients retained. Lambda changes both commutation structure and energy scales; comparison does not isolate noncommutation.',
            'spectator': 'gamma_8 is absent from both H0 and the perturbation; it commutes with either Hamiltonian and G8(t)=1. Active-seven and all-eight averages are both supplied.',
            'array_layout': 'times[k]; cases[c].traces[i][k] is G_(i+1)(times[k]); each mean is an array indexed by k',
            'mean_absolute_definition': 'Mean of abs(G_i), NOT abs(mean(G_i))',
            'rounding': 'Curve outputs rounded to 14 decimal places; validation uses unrounded calculation',
            'limitations': ['This one-sided beta=0 calculation illustrates cancellation in averages.',
                            'No size-winding, thermalization, teleportation, or gravitational conclusion follows.',
                            'No reproduction of finite-temperature experimental curves is claimed.'],
            'sources': [
                {'title': 'Kobrin, Schuster and Yao, Comment on Traversable wormhole dynamics on a quantum processor', 'date': '2023-02-15', 'url': 'https://arxiv.org/abs/2302.07897', 'use': 'H0 coefficients and normalization'},
                {'title': 'Jafferis et al., Comment on Comment on Traversable wormhole dynamics on a quantum processor', 'date': '2023-03-27', 'url': 'https://arxiv.org/abs/2303.15423', 'use': 'H0 and the 0.3 psi1 psi2 psi3 psi5 perturbation'},
            ],
            'generator': 'generate_majorana_data.py; Python standard library and NumPy only',
            'numpy_version': np.__version__,
        },
        'times': times.tolist(),
        'cases': cases,
        'validation': validations,
        'perturbation_commutators_frobenius': perturbation_commutators,
        'baseline_benchmarks': snapshots,
    }


def analytic_checks():
    singlet = np.array([0, 1, -1, 0], dtype=complex)/math.sqrt(2)
    bell = np.outer(singlet, singlet.conj())
    A = [Z, X]
    B = [(Z+X)/math.sqrt(2), (Z-X)/math.sqrt(2)]
    chsh = np.kron(A[0], B[0])+np.kron(A[0], B[1])+np.kron(A[1], B[0])-np.kron(A[1], B[1])
    errors = {'partial_transpose_spectrum': 0., 'local_marginals': 0., 'CHSH': 0., 'negativity': 0., 'joint_distribution': 0.}
    presets = []
    for p in sorted(set(list(np.arange(101)/100)+[1/3, 1/math.sqrt(2)])):
        rho = p*bell+(1-p)*np.eye(4)/4
        tensor = rho.reshape(2, 2, 2, 2)
        partial = tensor.transpose(0, 3, 2, 1).reshape(4, 4)
        eig = np.linalg.eigvalsh(partial)
        predicted = np.sort([(1-3*p)/4]+[(1+p)/4]*3)
        errors['partial_transpose_spectrum'] = max(errors['partial_transpose_spectrum'], maxabs(eig-predicted))
        errors['local_marginals'] = max(errors['local_marginals'],
            maxabs(np.trace(tensor, axis1=0, axis2=2)-I/2), maxabs(np.trace(tensor, axis1=1, axis2=3)-I/2))
        errors['CHSH'] = max(errors['CHSH'], abs(np.trace(rho@chsh).real+2*math.sqrt(2)*p))
        negativity = max(0., (3*p-1)/4)
        errors['negativity'] = max(errors['negativity'], abs(sum(-x for x in eig if x < 0)-negativity))
        for a in A:
            for b in B:
                dot = np.trace(a@b).real/2
                for s in (-1, 1):
                    for t in (-1, 1):
                        probability = np.trace(rho@np.kron((I+s*a)/2, (I+t*b)/2)).real
                        errors['joint_distribution'] = max(errors['joint_distribution'], abs(probability-(1-s*t*p*dot)/4))
        if p in (0., 1/3, .5, 1/math.sqrt(2), .85, 1.):
            spectrum = [(1+3*p)/4]+[(1-p)/4]*3
            entropy = -sum(x*math.log2(x) for x in spectrum if x > 0)
            presets.append({'p': p, 'CHSH_absolute': 2*math.sqrt(2)*p, 'negativity': negativity,
                            'mutual_information_bits': 2-entropy, 'partial_transpose_eigenvalues': predicted.tolist()})
    werner_validation = {name: check(name, error) for name, error in errors.items()}
    rt = []
    rt_error = 0.
    scale_error = 0.
    for ratio in (.1, .25, math.sqrt(2)-1, .5, 1.):
        length, cutoff, central_charge = 1., 1e-4, 30.
        gap = ratio*length
        sd = 2*central_charge/3*math.log(length/cutoff)
        sc = central_charge/3*math.log(gap*(2*length+gap)/cutoff**2)
        scaled_mi = max(0., math.log(1/(ratio*(2+ratio))))
        cross_ratio = 1/(1+ratio)**2
        other = max(0., math.log(cross_ratio/(1-cross_ratio)))
        rt_error = max(rt_error, abs(scaled_mi-other), abs(scaled_mi-3/central_charge*(sd-min(sd,sc))))
        factor = 7.3
        scaled_sd = 2*central_charge/3*math.log((factor*length)/(factor*cutoff))
        scaled_sc = central_charge/3*math.log((factor*gap)*(2*factor*length+factor*gap)/(factor*cutoff)**2)
        scale_error = max(scale_error, abs((sd-min(sd,sc))-(scaled_sd-min(scaled_sd,scaled_sc))))
        rt.append({'gap_over_length': ratio, 'cross_ratio': cross_ratio, 'three_MI_over_c': scaled_mi,
                   'S_disconnected': sd, 'S_connected': sc})
    return {
        'schema_version': 1,
        'status': 'Independent density-matrix and formula checks for browser analytic models; not empirical data',
        'werner': {'singlet': '(|01>-|10>)/sqrt(2)',
                   'thresholds': {'separable_through_p': 1/3, 'CHSH_violates_strictly_above_p': 1/math.sqrt(2)},
                   'threshold_semantics': 'Both inequalities are strict; equality is separable at 1/3 and CHSH-saturating at 1/sqrt(2). Classify from p rather than rounded labels.',
                   'settings': 'a0=z,a1=x,b0=(z+x)/sqrt(2),b1=(z-x)/sqrt(2); signed combination E00+E01+E10-E11 is negative',
                   'CHSH_scope': 'Nondegenerate spin measurements; no claim about every Bell scenario.',
                   'presets': presets, 'validation': werner_validation},
        'AdS3_two_intervals': {
            'assumptions': 'Vacuum on infinite line, equal intervals, static semiclassical AdS3, leading classical RT contribution, cutoff much smaller than gap and interval.',
            'units': 'Entropies in nats; table uses length=1, cutoff=0.0001, central charge=30.',
            'transition_gap_over_length': math.sqrt(2)-1,
            'transition_cross_ratio': .5,
            'formula': '3 I(A:B)/c = max(0, ln[1/(r(2+r))]), r=d/ell',
            'presets': rt,
            'validation': {'equivalent_formulas': check('RT equivalent formulas', rt_error),
                           'scale_invariance': check('RT scale invariance', scale_error)},
            'limitations': 'Connected entanglement wedge is not a traversable signal route. Zero leading-order mutual information need not be exact at finite c.',
        },
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output-dir', type=Path, default=Path(__file__).resolve().parent)
    args = parser.parse_args()
    majorana = majorana_data()
    analytic = analytic_checks()
    majorana['metadata']['generator_sha256'] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    for filename, payload in [('majorana-data.json', majorana), ('model-checks.json', analytic)]:
        path = args.output_dir / filename
        path.write_text(json.dumps(payload, indent=2, allow_nan=False)+'\n', encoding='utf-8')
        print(f'Wrote {path}')
    print('All checks passed; two cases, eight traces each, 241 time samples.')
    print(json.dumps({'baseline_t10': majorana['baseline_benchmarks'][2],
                      'RMS_energies': [case['rms_energy'] for case in majorana['cases']],
                      'validation': majorana['validation']}, indent=2))


if __name__ == '__main__':
    main()