#!/usr/bin/env python3
"""Check the OneScience/DTK TensorFlow runtime and build Saluki."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from _bootstrap import configure_runtime


ROOT = configure_runtime()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--params", type=Path, default=ROOT / "conf" / "params.json")
    parser.add_argument("--forward", action="store_true", help="Run one zero-input forward pass.")
    args = parser.parse_args()

    import numpy as np
    import tensorflow as tf

    from model.saluki import SalukiModel, load_params

    params = load_params(args.params)
    model = SalukiModel(params, head=0)
    report = {
        "tensorflow": tf.__version__,
        "built_with_gpu": tf.test.is_built_with_gpu_support(),
        "gpus": [device.name for device in tf.config.list_physical_devices("GPU")],
        "input_shape": list(model.keras_model.input_shape),
        "output_shape": list(model.keras_model.output_shape),
        "parameters": int(model.keras_model.count_params())
    }
    if args.forward:
        shape = model.keras_model.input_shape
        output = model.keras_model(
            np.zeros((1, shape[1], shape[2]), dtype=np.float32), training=False
        )
        output_np = output.numpy()
        report["forward"] = {
            "shape": list(output_np.shape),
            "finite": bool(np.isfinite(output_np).all()),
            "device": output.device
        }
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()