#!/usr/bin/env python3
"""Graviton Note 02: finite Majorana sensitivity sweep. Python 3.10+, NumPy.

Run beside protocol.json: python3 verify-robustness.py [--output-dir DIRECTORY]
Writes results.json and contrasts.csv. No network, random sampling, or measured
data. All results are predictions of the explicitly specified finite model.
"""
import argparse
import csv
import hashlib
import itertools
import json
import math
from functools import reduce
from pathlib import Path
import numpy as np

HERE = Path(__file__).resolve().parent
DIM = 16


def operators():
    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)
    return [reduce(np.kron, [Z]*j + [P] + [I]*(3-j))/np.sqrt(2)
            for j in range(4) for P in (X, Y)]


def center_rms(H):
    A = H - np.trace(H)/DIM*np.eye(DIM)
    return A, float(np.sqrt(np.trace(A@A).real/DIM))


def exponential(A):
    """Independent matrix exponential: Taylor expansion, scaling and squaring."""
    norm = float(np.linalg.norm(A, ord=np.inf))
    n = max(0, math.ceil(math.log2(norm))) if norm else 0
    B = A/(2**n)
    out = term = np.eye(DIM, dtype=complex)
    for k in range(1, 160):
        term = term@B/k
        out = out + term
        if np.linalg.norm(term, ord=np.inf) < 1e-17:
            break
    else:
        raise AssertionError('Matrix exponential failed to converge')
    for _ in range(n):
        out = out@out
    return out


def correlations(H, psi, times):
    E, V = np.linalg.eigh(H)
    gaps = (E[:, None]-E[None, :]).reshape(-1)
    weights = np.array([np.abs(V.conj().T@p@V).reshape(-1)**2 for p in psi])
    return (2/DIM)*weights@np.cos(gaps[:, None]*times[None, :]), E


def generate(protocol):
    psi = operators()
    identity = np.eye(DIM)
    quartic = lambda ids: reduce(np.matmul, [psi[i-1] for i in ids])
    terms = protocol['baseline_terms']
    used = {tuple(ids) for _, ids in terms}
    family = [ids for ids in itertools.combinations(range(1, 8), 4) if ids not in used]
    assert len(family) == 30
    H0 = sum(c*quartic(ids) for c, ids in terms)
    _, rms0 = center_rms(H0)
    fine_dt = min(protocol['refinement_steps'])
    times = np.arange(round(48/fine_dt)+1)*fine_dt
    base, _ = correlations(H0, psi, times)
    tol = protocol['verification_tolerance']
    errors = {}

    def check(name, error):
        val = float(error)
        errors[name] = max(errors.get(name, 0.), val)
        assert math.isfinite(val) and val <= tol, (name, val)

    check('majorana_algebra', max(np.max(np.abs(a@b+b@a-(identity if i==j else 0)))
          for i, a in enumerate(psi) for j, b in enumerate(psi)))
    check('baseline_terms_commute', max(np.max(np.abs(quartic(a)@quartic(b)-quartic(b)@quartic(a)))
          for _, a in terms for _, b in terms))
    check('zero_hamiltonian', np.max(np.abs(correlations(np.zeros_like(H0), psi, times[:5])[0]-1)))

    def scores(curves, step):
        stride = round(step/fine_dt)
        return [float(np.mean(np.abs(curves[:7, round(a/fine_dt):round(b/fine_dt)+1:stride])))
                for a, b in protocol['windows']]

    steps = [protocol['primary_step'], *protocol['refinement_steps']]
    base_scores = [scores(base, step) for step in steps]
    cases = []
    for ids in family:
        P = quartic(ids)
        commutator_norm = float(np.linalg.norm(H0@P-P@H0, ord='fro'))
        for lam in protocol['couplings']:
            raw, raw_rms = center_rms(H0+lam*P)
            H = raw*rms0/raw_rms
            curves, E = correlations(H, psi, times)
            check('hermiticity', np.max(np.abs(H-H.conj().T)))
            check('centered_trace', abs(np.trace(H)/DIM))
            check('rms_energy_match', abs(center_rms(H)[1]-rms0))
            check('eigenvalue_rms_match', abs(np.std(E)-rms0))
            check('initial_correlations', np.max(np.abs(curves[:, 0]-1)))
            check('spectator_conservation', np.max(np.abs(curves[7]-1)))
            check('absolute_bound_violation', max(0., float(np.max(np.abs(curves)))-1))
            for t in (0., 6., 12., 18., 24., 36., 48.):
                U = exponential(-1j*H*t)
                check('independent_unitarity', np.max(np.abs(U.conj().T@U-identity)))
                direct = np.array([(2/DIM)*np.trace(U.conj().T@p@U@p).real for p in psi])
                check('independent_exponential_vs_spectral', np.max(np.abs(direct-curves[:, round(t/fine_dt)])))
            ds = [scores(curves, step) for step in steps]
            contrasts = [[x-y for x, y in zip(d, b)] for d, b in zip(ds, base_scores)]
            cases.append({'indices':list(ids), 'label':''.join(map(str, ids)), 'lambda':lam,
                          'commutator_frobenius_norm':commutator_norm,
                          'rescaling_factor':rms0/raw_rms,
                          'spectral_width':float(E[-1]-E[0]),
                          'D_by_step':ds, 'contrasts_by_step':contrasts,
                          'mean_abs_trace':np.round(np.mean(np.abs(curves[:7, ::4]), axis=0), 12).tolist()})

    # Values recorded before this extension, rounded from the earlier dataset.
    earlier = json.loads((HERE.parents[1]/'inquiry/entanglement-spacetime/matched-test.json').read_text()) if (HERE.parents[1]/'inquiry/entanglement-spacetime/matched-test.json').exists() else None
    # Standalone reproduction also works without that optional historical dataset.
    if earlier:
        for prior in earlier['cases'][1:]:
            c = next(c for c in cases if c['label']=='1235' and c['lambda']==prior['perturbation_lambda'])
            check('historical_reference_reproduced', abs(c['contrasts_by_step'][0][1]-prior['contrast_D_minus_baseline']))
    counts=[]
    sign_tol=protocol['sign_tolerance']
    sign=lambda x: 1 if x>sign_tol else -1 if x < -sign_tol else 0
    for w, window in enumerate(protocol['windows']):
        vals=[c['contrasts_by_step'][0][w] for c in cases]
        counts.append({'window':window, 'negative':sum(sign(x)==-1 for x in vals),
                       'positive':sum(sign(x)==1 for x in vals), 'near_zero':sum(sign(x)==0 for x in vals)})
    refinements=[abs(c['contrasts_by_step'][2][w]-c['contrasts_by_step'][0][w]) for c in cases for w in range(4)]
    changes=[{'indices':c['indices'],'lambda':c['lambda'],'window':protocol['windows'][w]}
             for c in cases for w in range(4)
             if sign(c['contrasts_by_step'][0][w]) != sign(c['contrasts_by_step'][2][w])]
    flips=sum(min(c['contrasts_by_step'][0]) < -sign_tol and max(c['contrasts_by_step'][0]) > sign_tol for c in cases)
    summary={'cases':len(cases), 'window_comparisons':len(cases)*4,
             'counts_by_window':counts, 'cases_changing_sign_across_windows':flips,
             'max_primary_to_finest_contrast_change':max(refinements),
             'sign_changes_on_refinement':changes,
             'minimum_absolute_primary_contrast':min(abs(x) for c in cases for x in c['contrasts_by_step'][0]),
             'all_perturbations_commute_with_H0':all(c['commutator_frobenius_norm']<=tol for c in cases),
             'all_perturbations_noncommuting_with_H0':all(c['commutator_frobenius_norm']>tol for c in cases),
             'universal_negative_extension_passes':all(x < -sign_tol for c in cases for x in c['contrasts_by_step'][0]),
             'count_interpretation':'Finite enumeration only; not independent replications, probabilities, confidence intervals, or tests of spacetime geometry.'}
    return {'protocol':protocol, 'software':{'numpy':np.__version__},
            'baseline_rms_energy':rms0, 'steps':steps, 'baseline_D_by_step':base_scores,
            'plot_times':np.round(times[::4], 6).tolist(),
            'baseline_mean_abs_trace':np.round(np.mean(np.abs(base[:7, ::4]), axis=0), 12).tolist(),
            'summary':summary, 'validation':{'tolerance':tol,'max_errors':errors,'passed':True}, 'cases':cases}


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output-dir', type=Path, default=HERE)
    args=parser.parse_args()
    raw=(HERE/'protocol.json').read_bytes()
    result=generate(json.loads(raw))
    result['hashes']={'protocol_sha256':hashlib.sha256(raw).hexdigest(),
                      'generator_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
    args.output_dir.mkdir(parents=True, exist_ok=True)
    (args.output_dir/'results.json').write_text(json.dumps(result, separators=(',',':'), allow_nan=False)+'\n')
    with (args.output_dir/'contrasts.csv').open('w',newline='') as f:
        writer=csv.writer(f)
        writer.writerow(['quartic_indices','lambda','window_start','window_end','D_baseline_step_0.1','D_perturbed_step_0.1','contrast_step_0.1','contrast_step_0.05','contrast_step_0.025'])
        for c in result['cases']:
            for w, (a,b) in enumerate(result['protocol']['windows']):
                writer.writerow([c['label'],c['lambda'],a,b,result['baseline_D_by_step'][0][w],c['D_by_step'][0][w],*[v[w] for v in c['contrasts_by_step']]])
    print(json.dumps({'summary':result['summary'],'validation':result['validation'],'hashes':result['hashes']},indent=2))


if __name__=='__main__':
    main()