"""Offline graph/weight inspection; TensorFlow is confined to this tooling."""

import argparse
import hashlib
import json
import os
from pathlib import Path

os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

import numpy as np
import tensorflow as tf
from tensorflow.core.protobuf import saved_model_pb2


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=False)
    identity = json.loads((args.source / ".source.json").read_text())
    for name, expected in identity["files"].items():
        if hashlib.sha256((args.source / name).read_bytes()).hexdigest() != expected:
            raise ValueError("Source Google artifact changed: " + name)
    saved = saved_model_pb2.SavedModel()
    saved.ParseFromString((args.source / "saved_model.pb").read_bytes())
    graph = next(g for g in saved.meta_graphs if not g.meta_info_def.tags)
    (args.output / "graph.pb").write_bytes(graph.graph_def.SerializeToString())
    ops = []
    constants = {}
    for node in graph.graph_def.node:
        attrs = {}
        for name, value in node.attr.items():
            field = value.WhichOneof("value")
            if field in ("s", "i", "f", "b"):
                v = getattr(value, field)
                attrs[name] = v.decode(errors="replace") if isinstance(v, bytes) else v
            elif field == "list":
                attrs[name] = {k: list(getattr(value.list, k)) for k in ("i", "f", "b") if len(getattr(value.list, k))}
            elif field == "shape":
                attrs[name] = [d.size for d in value.shape.dim]
        ops.append({"name": node.name, "op": node.op, "inputs": list(node.input), "attrs": attrs})
        if node.op == "Const":
            array = tf.make_ndarray(node.attr["value"].tensor)
            if array.dtype.kind in "fiu" and array.size:
                constants[node.name] = array
    reader = tf.train.load_checkpoint(str(args.source / "variables/variables"))
    variables = {name: reader.get_tensor(name) for name in sorted(reader.get_variable_to_shape_map())}
    np.savez(args.output / "variables.npz", **variables)
    np.savez(args.output / "constants.npz", **constants)
    (args.output / "ops.json").write_text(json.dumps(ops, indent=2) + "\n")
    report = {"tensorflow": tf.__version__, "numpy": np.__version__, "source": identity,
              "tags": [list(m.meta_info_def.tags) for m in saved.meta_graphs],
              "variables": {k: list(v.shape) for k, v in variables.items()},
              "signatures": {name: {"inputs": {k: v.name for k, v in s.inputs.items()},
                                      "outputs": {k: v.name for k, v in s.outputs.items()}} for name, s in graph.signature_def.items()},
              "op_counts": {op: sum(n["op"] == op for n in ops) for op in sorted({n["op"] for n in ops})}}
    (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()