defmodule AxonInterp do
  @moduledoc """
  The thin wrapper for Axon inference.
  It is a Deep Learning inference framework that can be used in the same way as my *Interp.
  """

  @basic_usage """
  ## Basic Usage
  You get the trained axon model or onnx model and save it in a directory that your application can read.
  "your-app/priv" may be good choice.

  ```
  $ cp your-trained-model.axon ./priv
  ```

  Next, you will create a module that interfaces with the deep learning model.
  The module will need pre-processing and post-processing in addition to inference
  processing, as in the example following. AxonInterp provides inference processing
  only.

  You put `use AxonInterp` at the beginning of your module, specify the model path
  as an optional argument and describe model's inputs/outputs specification. In the
  inference, you will put data input to the model (`AxonInterp.set_input_tensor/3`),
  inference execution (`AxonInterp.invoke/1`), and inference result retrieval
  (`AxonInterp.get_output_tensor/2`).

  ```elixr:your_model.ex
  defmodule YourApp.YourModel do
    use AxonInterp,
      model: "priv/your-trained-model.onnx",
      inputs: [f32: {1, 3, 224, 224}],
      outputs: [f32: {1, 1000}]
      

    def predict(data) do
      # preprocess
      #  to convert the data to be inferred to the input format of the model.
      input_bin = convert-float32-binaries(data)

      # inference
      #  typical I/O data for Onnx models is a serialized 32-bit float tensor.
      output_bin = session()
        |> AxonInterp.set_input_tensor(0, input_bin)
        |> AxonInterp.invoke()
        |> AxonInterp.get_output_tensor(0)

      # postprocess
      #  add your post-processing here.
      #  you may need to reshape output_bin to tensor at first.
      tensor = output_bin
        |> Nx.from_binary({:f, 32})
        |> Nx.reshape({size-x, size-y, :auto})

      * your-postprocessing *
      ...
    end
  end
  ```
  """

  @timeout 300000

  @framework "axon"

  # the suffix expected for the model
  suffix = %{
    "axon" => [".axon", ".onnx"]
  }
  @model_suffix suffix[String.downcase(@framework)]
  
  # session record
  defstruct module: nil, inputs: %{}, outputs: %{}

  defmacro __using__(opts) do
    quote generated: true, location: :keep do
      use GenServer

      def start_link(opts) do
        GenServer.start_link(__MODULE__, opts, name: __MODULE__)
      end

      def init(opts) do
        opts = Keyword.merge(unquote(opts), opts)
        nn_model   = AxonInterp.validate_model(Keyword.get(opts, :model), Keyword.get(opts, :url))
        nn_inputs  = Keyword.get(opts, :inputs, [])
        nn_outputs = Keyword.get(opts, :outputs, [])

        {model, params} = case Path.extname(nn_model) do
          ".axon" -> File.read!(nn_model) |> Axon.deserialize()
          ".onnx" -> AxonOnnx.import(nn_model)
        end
        {_, predict_fn} = Axon.build(model, mode: :inference)

        {:ok, %{model: predict_fn, params: params, path: nn_model, itempl: nn_inputs, otempl: nn_outputs}}
      end

      def session() do
        %AxonInterp{module: __MODULE__}
      end

      def handle_call({:info}, _from, state) do
        info = %{
          "model"   => state.path,
          "inputs"  => state.itempl,
          "outputs" => state.otempl,
        }
        {:reply, {:ok, info}, state}
      end

      require Nx

      def handle_call({:invoke, inputs}, _from, %{model: model, params: params, itempl: template}=state) do
        n = Enum.count(template)

        inputs =
          if n == 1 do
            inputs[0]
          else
            Enum.map(0..n-1, &inputs[&1]) |> List.to_tuple()
          end

        results = case model.(params, inputs) do
          single when Nx.is_tensor(single) ->
            %{0 => single}
          multiple when is_tuple(multiple) ->
            Tuple.to_list(multiple)
            |> Enum.with_index()
            |> Enum.into(%{}, fn {result, index} -> {index, result} end)
        end
        
        {:reply, {:ok, results}, state}
      end

      def handle_call({:itempl, index}, _from, %{itempl: template}=state) do
        {:reply, {:ok, Enum.at(template, index)}, state}
      end

      def handle_call({:otempl, index}, _from, %{otempl: template}=state) do
        {:reply, {:ok, Enum.at(template, index)}, state}
      end

      def terminate(_reason, state) do
        :ok
      end
    end
  end


  @doc """
  Get name of backend NN framework.
  """
  def framework() do
    @framework
  end

  @doc """
  Ensure that the back-end framework is as expected.
  """
  def framework?(name) do
    unless String.downcase(name) == String.downcase(@framework),
      do: raise "Error: backend NN framework is \"#{@framework}\", not \"#{name}\"."
  end

  @doc """
  Ensure that the model matches the back-end framework.
  
  ## Parameters
    * model - path of model file
    * url - download site
  """
  def validate_model(nil, _), do: raise ArgumentError, "need a model file #{inspect(@model_suffix)}."
  def validate_model(model, url) do
    validate_extname!(model)

    abs_path = Path.expand(model)
    unless File.exists?(abs_path) do
      IO.puts("#{model}:")
      {:ok, _} = AxonInterp.URL.download(url, Path.dirname(abs_path), Path.basename(abs_path))
    end
    model
  end

  defp validate_extname!(model) do
    actual_ext = Path.extname(model)
    unless actual_ext in @model_suffix,
      do: raise ArgumentError, "expects the model file #{inspect(@model_suffix)} not \"#{actual_ext}\"."

    actual_ext
  end

  @doc """
  Get the propaty of the model.

  ## Parameters

    * mod - modules' names
  """
  def info(mod) do
    case GenServer.call(mod, {:info}, @timeout) do
      {:ok, result} ->  {:ok, Map.put(result, "framework", @framework)}
      any -> any
    end
  end

  @doc """
  Stop the axon interpreter.

  ## Parameters

    * mod - modules' names
  """
  def stop(mod) do
    GenServer.stop(mod)
  end

  @doc """
  Put a flat binary to the input tensor on the interpreter.

  ## Parameters

    * sessin - session record
    * index - index of input tensor in the model
    * bin   - input data - flat binary, cf. serialized tensor
    * opts
      * [:binary] - data is binary
  """
  def set_input_tensor(%AxonInterp{module: mod, inputs: inputs}=session, index, data, opts \\ []) do
    data = cond do
      :binary in opts ->
        with {:ok, {dtype, shape}} <- GenServer.call(mod, {:itempl, index}, @timeout) do
          Nx.from_binary(data, dtype) |> Nx.reshape(shape)
        end

      true -> data
      end

    %AxonInterp{session | inputs: Map.put(inputs, index, data)}
  end

  @doc """
  Get the flat binary from the output tensor on the interpreter.

  ## Parameters

    * session - session record
    * index - index of output tensor in the model
    * opts:
      * [:binary] - convert the output to binary
  """
  def get_output_tensor(%AxonInterp{outputs: outputs}, index, opts \\ []) do
    output = outputs[index]

    cond do
      :binary in opts -> Nx.to_binary(output)
      true -> output
    end
  end

  @doc """
  Execute the inference session. In session mode, data input/execution of
  inference/output of results to the DL model is done all at once.

  ## Parameters

    * session - session

  ## Examples.

    ```elixir
      output_bin0 = session()
        |> AxonInterp.set_input_tensor(0, input_bin0)
        |> AxonInterp.invoke()
        |> AxonInterp.get_output_tensor(0)
    ```
  """
  def invoke(%AxonInterp{module: mod, inputs: inputs}=session) do
    case GenServer.call(mod, {:invoke, inputs}, @timeout) do
      {:ok, results} -> %AxonInterp{session | outputs: results}
      any -> any
    end
  end
end