open Js_of_ocaml
open Onnxrt

let log msg =
  let el = Dom_html.getElementById "output" in
  let prev =
    Js.to_string (Js.Opt.get el##.textContent (fun () -> Js.string ""))
  in
  el##.textContent := Js.some (Js.string (prev ^ msg ^ "\n"))

let () =
  Lwt.async @@ fun () ->
  let open Lwt.Syntax in
  log "Creating tensors...";
  let a_data =
    Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout
      [| 1.0; 2.0; 3.0 |]
  in
  let b_data =
    Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout
      [| 4.0; 5.0; 6.0 |]
  in
  let a = Tensor.of_bigarray1 Dtype.Float32 a_data ~dims:[| 3 |] in
  let b = Tensor.of_bigarray1 Dtype.Float32 b_data ~dims:[| 3 |] in
  log "Loading model...";
  let* session = Session.create "add.onnx" () in
  log "Running inference...";
  let* outputs = Session.run session [ ("A", a); ("B", b) ] in
  let c = List.assoc "C" outputs in
  let c_data = Tensor.to_bigarray1_exn Dtype.Float32 c in
  log
    (Printf.sprintf "Result C = [%g, %g, %g]" (Bigarray.Array1.get c_data 0)
       (Bigarray.Array1.get c_data 1)
       (Bigarray.Array1.get c_data 2));
  Tensor.dispose a;
  Tensor.dispose b;
  Tensor.dispose c;
  let* () = Session.release session in
  log "Done!";
  Lwt.return_unit