import json
import numpy as np
from sparse_lyapunov import rho_scalar, run_experiment, verify_math


def summarize_trials(seeds=range(10)):
    rows = [run_experiment(seed=int(s), steps=180) for s in seeds]
    b20 = np.array([r['baseline']['loss_20pct'] for r in rows])
    i20 = np.array([r['idea']['loss_20pct'] for r in rows])
    bf = np.array([r['baseline']['final_loss'] for r in rows])
    inf = np.array([r['idea']['final_loss'] for r in rows])
    bs = np.array([r['baseline']['spikes'] for r in rows])
    ins = np.array([r['idea']['spikes'] for r in rows])
    checked = []
    for r in rows:
        for cert, curv in zip(r['idea']['certificates'], r['curvatures']):
            checked.append(max(rho_scalar(cert['eta'], cert['beta'], float(v)) for v in curv))
    return {
        'math': verify_math(), 'n_trials': len(rows),
        'baseline_mean_loss_20pct': float(b20.mean()),
        'idea_mean_loss_20pct': float(i20.mean()),
        'baseline_mean_final_loss': float(bf.mean()),
        'idea_mean_final_loss': float(inf.mean()),
        'baseline_mean_spikes': float(bs.mean()),
        'idea_mean_spikes': float(ins.mean()),
        'early_loss_ratio_idea_over_baseline': float(np.mean(i20 / b20)),
        'all_reported_endpoint_rhos_lt_1': bool(np.all(np.array(checked) < 1.0)),
        'max_checked_endpoint_rho': float(max(checked)),
        'per_seed': [{'seed': int(s), 'baseline_loss_20pct': float(a),
                      'idea_loss_20pct': float(b), 'baseline_spikes': int(c),
                      'idea_spikes': int(d)} for s,a,b,c,d in zip(seeds,b20,i20,bs,ins)]
    }

if __name__ == '__main__':
    print(json.dumps(summarize_trials(), indent=2))