#!/usr/bin/env python3
"""Write web/data.json: the few things the browser demos take straight from the
Python scripts, so the page uses exactly the same inputs.

  * double descent: the 12 training inputs x and the unit noise draws z
    (make_data(seed=0) gives y = sin(2 pi x) + 0.1 z);
  * certified robustness: the 200 test points and the weights of the trained
    wiggly and Lipschitz networks (build_models(quick=False), about 20 s).

  python web/export_data.py
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(HERE))

import certified_robustness as cr  # noqa: E402
import double_descent as dd  # noqa: E402


def r(a, k=8):
    return np.round(np.asarray(a, float), k).tolist()


def net_json(net):
    return {"W": [r(W) for W in net.W], "b": [r(b) for b in net.b],
            "L": net.lipschitz_bound()}


def main():
    x, y = dd.make_data()
    z = (y - dd.target(x)) / dd.NOISE
    (_, _, Xte, yte), wiggly, lip = cr.build_models(quick=False)
    data = {
        "dd": {"x": r(x, 12), "z": r(z, 12), "noise": dd.NOISE, "scale": dd.SCALE, "p_grid": dd.P_GRID},
        "cr": {"Xte": r(Xte, 6), "yte": yte.astype(int).tolist(),
               "wiggly": net_json(wiggly), "lip": net_json(lip)},
    }
    path = os.path.join(HERE, "data.json")
    with open(path, "w") as fh:
        json.dump(data, fh, separators=(",", ":"))
    print(f"wrote {path} ({os.path.getsize(path) // 1024} KB)")


if __name__ == "__main__":
    main()