import json
import numpy as np
from teleport_simplicial import (
    cyclic_triangle_incidence, local_operator, teleport, nonprincipal_radius,
    mfpt_to_target, run_experiment,
)


def main():
    out = run_experiment(seed=7)
    n = out['n']
    A = cyclic_triangle_incidence(n)
    P0 = local_operator(A, self_loop=0.1)
    alphas = np.array([x[0] for x in out['spectral']])

    # Structural checks: incidence has three faces per triangle and P is stochastic.
    assert np.all(A.sum(axis=0) == 3)
    assert np.allclose(P0.sum(axis=1), 1.0)
    assert np.min(P0) >= 0

    # Prediction 1: uniform teleportation gives rho(P_alpha)=rho(P0)*(1-alpha).
    rho0 = out['rho0']
    observed = np.array([x[1] for x in out['spectral']])
    predicted = rho0 * (1 - alphas)
    spectral_max_error = float(np.max(np.abs(observed - predicted)))
    assert spectral_max_error < 1e-10

    # Prediction 2: global teleportation monotonically reduces target MFPT.
    mfpt = np.array([x[1] for x in out['mfpt']])
    assert np.all(np.diff(mfpt) < 0)

    # A direct mixing prediction: on any mean-zero eigenmode, m-step norm scales
    # as [rho(P0)*(1-alpha)]^m.  Numerically use the slowest eigenvector.
    vals, vecs = np.linalg.eig(P0)
    k = np.argsort(np.abs(vals - 1))[-1]  # not used; select largest nonprincipal below
    order = np.argsort(-np.abs(vals))
    k = next(i for i in order if abs(vals[i] - 1) > 1e-8)
    v = np.real(vecs[:, k]); v -= v.mean(); v /= np.linalg.norm(v)
    m = 12
    mixing_rows = []
    for a in alphas:
        P = teleport(P0, float(a))
        actual = np.linalg.norm(np.linalg.matrix_power(P, m) @ v)
        pred = (rho0 * (1-a)) ** m
        mixing_rows.append([float(a), float(actual), float(pred)])
    mixing_max_error = float(max(abs(x[1]-x[2]) for x in mixing_rows))
    assert mixing_max_error < 1e-8

    report = {
        'spectral_prediction': {
            'prediction': 'rho_alpha / rho_0 = 1-alpha',
            'max_abs_error': spectral_max_error,
            'rows_alpha_observed_predicted': [
                [float(a), float(o), float(p)] for a, o, p in zip(alphas, observed, predicted)
            ],
        },
        'mfpt_prediction': {
            'prediction': 'mean MFPT to target decreases with alpha',
            'rows_alpha_mean_mfpt': [[float(a), float(t)] for a, t in zip(alphas, mfpt)],
        },
        'mixing_prediction': {
            'prediction': 'm-step slow-mode norm = [rho_0(1-alpha)]^m, m=12',
            'max_abs_error': mixing_max_error,
            'rows_alpha_observed_predicted': mixing_rows,
        },
        'classification': {
            'prediction': 'teleportation can improve delayed global classification but oversmooths features',
            'rows_alpha_accuracy_feature_std': out['classification'],
        },
    }
    with open('results.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()