"""The cost of instantiation against the cost of the network.

One generation of the host is: inst splits the emitted bytes, sigma^-1 parses
the first field into the netlist, Lev places the rows and entries of the
layered map, and the device is applied once per step. This script counts those
operations, one per byte read by inst and by sigma^-1, one per row and one per
nonzero entry placed by Lev, and one per device application, times each of
them, and compares the total with the T * M unit evaluations of the run.
"""
import json
import os
import sys
import time

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

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def main() -> int:
    from selfrep import (M_STAR, inst, ser, tau_star, run_reference, read_host,
                         host_netlist, sigma_host)
    from netlist_io import net_of_sigma
    from matrix8 import compile_net

    sigma_bytes = read_host()
    tau = tau_star(sigma_bytes, M_STAR)
    s = ser(sigma_bytes, bytes(M_STAR), tau)

    # sigma: parameters and predecessor entries stored, and the time to write
    t0 = time.perf_counter()
    net, inputs, outputs = host_netlist()
    t_build = time.perf_counter() - t0
    t0 = time.perf_counter()
    encoded = sigma_host()
    t_sigma = time.perf_counter() - t0
    assert encoded == sigma_bytes
    stored, _, _, _ = net_of_sigma(sigma_bytes)
    entries = (2 * len(stored.gates)                       # bias and fan-in
               + sum(len(i) for i, _ in stored.gates.values())
               + len(outputs))
    print(f"  sigma: {entries:,} stored entries -> {len(encoded):,} bytes "
          f"in {t_sigma * 1e3:.0f} ms ({t_build * 1e3:.0f} ms to assemble "
          f"the netlist from source)")

    # inst: one pass over the serialized instance, then sigma^-1 on its first field
    t0 = time.perf_counter()
    reps = 50
    for _ in range(reps):
        a, b, c = inst(s)
    t_inst = (time.perf_counter() - t0) / reps
    assert (a, b, c) == (sigma_bytes, bytes(M_STAR), tau)
    t0 = time.perf_counter()
    for _ in range(5):
        N, I, O, _ = net_of_sigma(a)
    t_sigma_inv = (time.perf_counter() - t0) / 5
    assert len(N.gates) == len(stored.gates)
    print(f"  inst: {len(s):,} bytes parsed in {t_inst * 1e3:.2f} ms, and "
          f"{len(N.gates):,} units rebuilt from the first field in "
          f"{t_sigma_inv * 1e3:.0f} ms")

    # Lev: the units placed, and the time to place them
    t0 = time.perf_counter()
    layers, info = compile_net(N, I, O)
    t_lev = time.perf_counter() - t0
    M = sum(int(W.shape[0]) for W, _ in layers)
    nnz = sum(int((W != 0).sum()) for W, _ in layers)
    print(f"  Lev: {len(N.gates):,} units in, {M:,} rows and {nnz:,} nonzero "
          f"entries placed over {info['layers']} layers in {t_lev:.1f} s")

    # the device: one application per step, at most one byte moved
    t0 = time.perf_counter()
    emitted, T = run_reference(M_STAR, tau)
    t_run = time.perf_counter() - t0
    assert emitted == s
    print(f"  device: {T:,} applications, {len(emitted):,} bytes emitted "
          f"({t_run:.2f} s for the whole reference run)")

    env_ops = len(s) + len(sigma_bytes) + M + nnz + T
    closed = 2 * len(sigma_bytes) + len(bytes(M_STAR)) + len(tau) + 16 + M + nnz + T
    assert env_ops == closed, (env_ops, closed)
    net_ops = T * M
    print(f"  instantiation: {env_ops:,} elementary operations per generation, "
          f"which is the closed form 2|sigma|+|m|+|tau|+16+M+nnz+T")
    print(f"  network:     {net_ops:,} unit evaluations (T = {T:,}, M = {M:,})")
    print(f"  ratio: {env_ops / net_ops:.3e}")

    out = {"sigma_entries": entries, "sigma_bytes": len(encoded),
           "sigma_seconds": t_sigma, "build_seconds": t_build,
           "inst_bytes": len(s), "inst_seconds": t_inst,
           "sigma_inverse_seconds": t_sigma_inv,
           "lev_units_in": len(N.gates), "lev_units_out": M, "lev_nonzero": nnz,
           "lev_layers": info["layers"], "lev_seconds": t_lev,
           "lev_entries": info["total_weights"],
           "steps": T, "emitted_bytes": len(emitted),
           "reference_seconds": t_run,
           "environment_ops": env_ops, "network_ops": net_ops,
           "ratio": env_ops / net_ops}
    json.dump(out, open(runs_path("paper_environment.json"), "w"), indent=1)
    return 0


if __name__ == "__main__":
    sys.exit(main())