cix_p1_tpu

defmodule CixP1 do
  @moduledoc """
  Elixir bindings for the CIX P1 (Arm-China "Zhouyi") NPU on the Orange Pi 6 / 6
  Plus, via the CIX NOE runtime (`libnoe`).

  This is a thin, low-level wrapper over the NOE inference API. It loads a
  pre-compiled graph binary and runs inference on raw tensor binaries; model
  pre/post-processing (image resize, detection decoding, …) is left to the
  caller. An optional `CixP1.Nx` layer converts tensor binaries to/from
  `Nx.Tensor` when `:nx` is available.

  ## Requirements

  Runs only on a Nerves image built from `nerves_system_orangepi6` (v0.2.0+),
  which provides the `aipu` kernel driver, `libnoe`/`libaipudrv`, and the
  `LD_LIBRARY_PATH` needed to resolve them. Models must be compiled to the
  NOE `.cix` format (NOE Compiler / CIX AI Model Hub).

  ## One-shot inference

      {:ok, ctx}  = CixP1.Context.new()
      {:ok, graph} = CixP1.Graph.load(ctx, "/data/models/mobilenet.cix")
      {:ok, [logits]} = CixP1.run(graph, [image_binary])

  See `CixP1.Context`, `CixP1.Graph`, and `CixP1.Job` for the step-by-step API.
  """

  alias CixP1.{Graph, Job}

  @doc """
  Runs one inference: creates a job on `graph`, loads `inputs` (a list of raw
  binaries, one per input tensor in order), runs the NPU, and collects every
  output tensor as a raw binary.

  `opts`:
    * `:timeout_ms` — inference timeout (default `5000`; `<= 0` waits forever)

  Returns `{:ok, [binary]}` (outputs in tensor order) or `{:error, reason}`.
  """
  @spec run(Graph.t(), [binary()], keyword()) ::
          {:ok, [binary()]} | {:error, String.t()}
  def run(%Graph{} = graph, inputs, opts \\ []) when is_list(inputs) do
    timeout = Keyword.get(opts, :timeout_ms, 5_000)

    with {:ok, job} <- Job.create(graph),
         :ok <- load_inputs(job, inputs),
         :ok <- Job.infer(job, timeout),
         {:ok, count} <- Graph.output_count(graph),
         {:ok, outputs} <- collect_outputs(job, count) do
      {:ok, outputs}
    end
  end

  defp load_inputs(job, inputs) do
    inputs
    |> Enum.with_index()
    |> Enum.reduce_while(:ok, fn {data, index}, :ok ->
      case Job.load_input(job, index, data) do
        :ok -> {:cont, :ok}
        {:error, _} = err -> {:halt, err}
      end
    end)
  end

  defp collect_outputs(_job, 0), do: {:ok, []}

  defp collect_outputs(job, count) do
    Enum.reduce_while(0..(count - 1)//1, {:ok, []}, fn i, {:ok, acc} ->
      case Job.get_output(job, i) do
        {:ok, data} -> {:cont, {:ok, [data | acc]}}
        {:error, _} = err -> {:halt, err}
      end
    end)
    |> case do
      {:ok, acc} -> {:ok, Enum.reverse(acc)}
      err -> err
    end
  end
end