cix_p1_tpu

defmodule CixP1.Graph do
  @moduledoc """
  A compiled model graph loaded onto the NPU.

  Load an offline-built graph binary (`.cix`, produced by the NOE Compiler or
  pulled from the CIX AI Model Hub) with `load/2`, then inspect its input and
  output tensors via the descriptor functions. Run inference through
  `CixP1.Job` or the `CixP1.run/2` convenience.
  """

  alias CixP1.{Context, Nif}

  @enforce_keys [:resource, :context]
  defstruct [:resource, :context]

  @type t :: %__MODULE__{resource: reference(), context: Context.t()}

  @typedoc """
  A tensor descriptor as reported by libnoe.

  * `:id` — tensor id within the graph
  * `:size` — buffer size in bytes
  * `:scale` / `:zero_point` — quantization parameters (`real = (q - zero_point) * scale`)
  * `:data_type` — element type atom (`:u8`, `:s8`, `:s16`, `:s32`, `:f16`, `:f32`, …)
    or the raw integer for types without an atom mapping
  """
  @type descriptor :: %{
          id: non_neg_integer(),
          size: non_neg_integer(),
          scale: float(),
          zero_point: integer(),
          data_type: atom() | non_neg_integer()
        }

  @doc """
  Loads a compiled graph binary from `path` into `context`.
  """
  @spec load(Context.t(), Path.t()) :: {:ok, t()} | {:error, String.t()}
  def load(%Context{resource: ctx} = context, path) when is_binary(path) do
    case Nif.graph_load(ctx, path) do
      {:ok, resource} -> {:ok, %__MODULE__{resource: resource, context: context}}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  @doc "Number of input tensors the graph expects."
  @spec input_count(t()) :: {:ok, non_neg_integer()} | {:error, String.t()}
  def input_count(%__MODULE__{resource: g}), do: count(g, Nif.tensor_type_input())

  @doc "Number of output tensors the graph produces."
  @spec output_count(t()) :: {:ok, non_neg_integer()} | {:error, String.t()}
  def output_count(%__MODULE__{resource: g}), do: count(g, Nif.tensor_type_output())

  @doc "Descriptor for input tensor `index` (0-based)."
  @spec input_descriptor(t(), non_neg_integer()) ::
          {:ok, descriptor()} | {:error, String.t()}
  def input_descriptor(%__MODULE__{resource: g}, index),
    do: descriptor(g, Nif.tensor_type_input(), index)

  @doc "Descriptor for output tensor `index` (0-based)."
  @spec output_descriptor(t(), non_neg_integer()) ::
          {:ok, descriptor()} | {:error, String.t()}
  def output_descriptor(%__MODULE__{resource: g}, index),
    do: descriptor(g, Nif.tensor_type_output(), index)

  @doc "All input descriptors, in order."
  @spec input_descriptors(t()) :: {:ok, [descriptor()]} | {:error, String.t()}
  def input_descriptors(%__MODULE__{} = graph),
    do: descriptors(graph, &input_count/1, &input_descriptor/2)

  @doc "All output descriptors, in order."
  @spec output_descriptors(t()) :: {:ok, [descriptor()]} | {:error, String.t()}
  def output_descriptors(%__MODULE__{} = graph),
    do: descriptors(graph, &output_count/1, &output_descriptor/2)

  defp count(g, type) do
    case Nif.graph_tensor_count(g, type) do
      {:ok, cnt} -> {:ok, cnt}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  defp descriptor(g, type, index) do
    case Nif.graph_tensor_descriptor(g, type, index) do
      {:ok, desc} -> {:ok, desc}
      {:error, reason} -> {:error, to_string(reason)}
    end
  end

  defp descriptors(graph, count_fun, desc_fun) do
    with {:ok, n} <- count_fun.(graph) do
      Enum.reduce_while(0..(n - 1)//1, {:ok, []}, fn i, {:ok, acc} ->
        case desc_fun.(graph, i) do
          {:ok, desc} -> {:cont, {:ok, [desc | acc]}}
          {:error, _} = err -> {:halt, err}
        end
      end)
      |> case do
        {:ok, acc} -> {:ok, Enum.reverse(acc)}
        err -> err
      end
    end
  end
end