import json
import numpy as np
from experiment import make_data, graph_from_x, cut_ratio, train_once, SEED

def run(seed):
    X, y = make_data(seed)
    rng = np.random.RandomState(seed + 101)
    policies = [X, X + rng.randn(*X.shape).astype('float32') * 0.18,
                X[rng.permutation(len(X))]]
    graphs = [graph_from_x(z, k=10, sigma=0.8) for z in policies]
    cuts = [cut_ratio(W, y) for W in graphs]
    uniform = train_once('uniform', X, y, graphs, cuts, seed + 17)
    cut = train_once('cut', X, y, graphs, cuts, seed + 17)
    return {
        'seed': seed, 'cuts': cuts,
        'uniform_accuracy': uniform['accuracy'],
        'cut_accuracy': cut['accuracy'],
        'uniform_pred_cut': uniform['pred_cut_uniform_graph'],
        'cut_pred_cut': cut['pred_cut_uniform_graph'],
        'cut_q': cut['q']
    }

rows = [run(s) for s in [SEED, SEED + 1, SEED + 2]]
out = {'runs': rows}
for key in ['uniform_accuracy', 'cut_accuracy', 'uniform_pred_cut', 'cut_pred_cut']:
    values = np.array([r[key] for r in rows])
    out[key + '_mean'] = float(values.mean())
    out[key + '_std'] = float(values.std())
print(json.dumps(out, indent=2))
with open('robustness.json', 'w') as f:
    json.dump(out, f, indent=2)