defmodule Mix.Tasks.DetectImage do
  @moduledoc """
  Image detection mix task: `mix help detect_image`

  Command line arguments:

  - `-m`, `--model`: *Required*. File path of .tflite file.
  - `-i`, `--input`: *Required*. Image to process.
  - `-l`, `--labels`: File path of labels file.
  - `-t`, `--threshold`: Default to `0.4`. Score threshold for detected objects.
  - `-c`, `--count`: Default to `1`. Number of times to run inference.
  - `-j`, `--jobs`: Number of threads for the interpreter (only valid for CPU).

  Code based on [detect_image.py](https://github.com/google-coral/pycoral/blob/master/examples/detect_image.py)
  """

  use Mix.Task

  alias TFLiteBEAM.Interpreter
  alias TFLiteBEAM.InterpreterBuilder
  alias TFLiteBEAM.TFLiteTensor
  alias TFLiteBEAM.FlatBufferModel

  @shortdoc "Object Detection"
  def run(argv) do
    {args, _, _} =
      OptionParser.parse(argv,
        strict: [
          model: :string,
          input: :string,
          labels: :string,
          threshold: :float,
          count: :integer,
          jobs: :integer
        ],
        aliases: [
          m: :model,
          i: :input,
          l: :labels,
          t: :threshold,
          c: :count,
          j: :jobs
        ]
      )

    default_values = [threshold: 0.4, count: 1, jobs: System.schedulers_online()]

    args =
      Keyword.merge(args, default_values, fn _k, user, default ->
        if user == nil do
          default
        else
          user
        end
      end)

    model = load_model(args[:model])
    input_image = %StbImage{shape: {h, w, _c}} = load_input(args[:input])
    labels = load_labels(args[:labels])
    interpreter = make_interpreter(model, args[:jobs])
    :ok = Interpreter.allocate_tensors(interpreter)

    [input_tensor_number | _] = Interpreter.inputs!(interpreter)
    output_tensor_numbers = Interpreter.outputs!(interpreter)

    if Enum.count(output_tensor_numbers) != 4 do
      raise ArgumentError, "Object detection models should have 4 output tensors"
    end

    %TFLiteTensor{} = input_tensor = Interpreter.tensor(interpreter, input_tensor_number)

    if input_tensor.type != {:u, 8} do
      raise ArgumentError, "Only support uint8 input type."
    end

    {height, width} =
      case input_tensor.shape do
        {_n, height, width, _c} ->
          {height, width}

        {_n, height, width} ->
          {height, width}

        shape ->
          raise RuntimeError, "not sure the input shape, got #{inspect(shape)}"
      end

    scale = min(height / h, width / w)
    {h, w} = {trunc(h * scale), trunc(w * scale)}

    input_image =
      StbImage.resize(input_image, h, w)
      |> StbImage.to_nx()
      |> Nx.new_axis(0)

    Nx.broadcast(0, input_tensor.shape)
    |> Nx.as_type(:u8)
    |> Nx.put_slice([0, 0, 0, 0], input_image)
    |> Nx.to_binary()
    |> then(&TFLiteTensor.set_data(input_tensor, &1))

    IO.puts("----INFERENCE TIME----")

    for _ <- 1..args[:count] do
      start_time = :os.system_time(:microsecond)
      Interpreter.invoke!(interpreter)
      end_time = :os.system_time(:microsecond)
      inference_time = (end_time - start_time) / 1000.0
      IO.puts("#{Float.round(inference_time, 1)}ms")
    end

    signature_list = Interpreter.get_signature_defs!(interpreter)

    {count_tensor_id, scores_tensor_id, class_ids_tensor_id, boxes_tensor_id} =
      if signature_list != nil do
        signature_list = Map.values(signature_list)

        if Enum.count(signature_list) > 1 do
          raise ArgumentError, "Only support model with one signature."
        else
          count_tensor_id = signature_list[:outputs][:output_0]
          scores_tensor_id = signature_list[:outputs][:output_1]
          class_ids_tensor_id = signature_list[:outputs][:output_2]
          boxes_tensor_id = signature_list[:outputs][:output_3]
          {count_tensor_id, scores_tensor_id, class_ids_tensor_id, boxes_tensor_id}
        end
      else
        %TFLiteTensor{} =
          output_tensor_3 = Interpreter.tensor(interpreter, Enum.at(output_tensor_numbers, 3))

        if output_tensor_3.shape == {1} do
          boxes_tensor_id = Enum.at(output_tensor_numbers, 0)
          class_ids_tensor_id = Enum.at(output_tensor_numbers, 1)
          scores_tensor_id = Enum.at(output_tensor_numbers, 2)
          count_tensor_id = Enum.at(output_tensor_numbers, 3)
          {count_tensor_id, scores_tensor_id, class_ids_tensor_id, boxes_tensor_id}
        else
          boxes_tensor_id = Enum.at(output_tensor_numbers, 1)
          class_ids_tensor_id = Enum.at(output_tensor_numbers, 3)
          scores_tensor_id = Enum.at(output_tensor_numbers, 0)
          count_tensor_id = Enum.at(output_tensor_numbers, 2)
          {count_tensor_id, scores_tensor_id, class_ids_tensor_id, boxes_tensor_id}
        end
      end

    boxes =
      Interpreter.tensor(interpreter, boxes_tensor_id)
      |> TFLiteTensor.to_nx(backend: Nx.BinaryBackend)
      |> take_first_and_reshape()

    class_ids =
      Interpreter.tensor(interpreter, class_ids_tensor_id)
      |> TFLiteTensor.to_nx(backend: Nx.BinaryBackend)
      |> take_first_and_reshape()

    scores =
      Interpreter.tensor(interpreter, scores_tensor_id)
      |> TFLiteTensor.to_nx(backend: Nx.BinaryBackend)
      |> take_first_and_reshape()

    count =
      Interpreter.tensor(interpreter, count_tensor_id)
      |> TFLiteTensor.to_nx(backend: Nx.BinaryBackend)
      |> Nx.to_flat_list()
      |> hd()
      |> trunc()

    {sx, sy} = {height / scale, width / scale}

    Enum.each(0..(count - 1), fn index ->
      score =
        Nx.take(scores, index)
        |> Nx.to_flat_list()
        |> hd()

      if score >= args[:threshold] do
        [ymin, xmin, ymax, xmax] =
          Nx.take(boxes, index)
          |> Nx.multiply(scale)
          |> Nx.to_flat_list()

        {xmin, xmax} = {trunc(sx * xmin), trunc(sx * xmax)}
        {ymin, ymax} = {trunc(sy * ymin), trunc(sy * ymax)}

        class_id =
          Nx.take(class_ids, index)
          |> Nx.to_flat_list()
          |> hd()
          |> trunc()

        class_str =
          if labels != nil do
            Enum.at(labels, class_id)
          else
            class_id
          end

        IO.puts("#{class_str}")
        IO.puts("  id   : #{class_id}")
        IO.puts("  score: #{Float.round(score, 3)}")
        IO.puts("  bbox : #{inspect([ymin, xmin, ymax, xmax])}")
      end
    end)
  end

  defp load_model(nil) do
    raise ArgumentError, "empty value for argument '--model'"
  end

  defp load_model(model_path) do
    FlatBufferModel.build_from_buffer(File.read!(model_path))
  end

  defp load_input(nil) do
    raise ArgumentError, "empty value for argument '--input'"
  end

  defp load_input(input_path) do
    with {:ok, input_image} <- StbImage.read_file(input_path) do
      input_image
    else
      {:error, error} ->
        raise RuntimeError, error
    end
  end

  defp load_labels(nil), do: nil

  defp load_labels(label_file_path) do
    File.read!(label_file_path)
    |> String.split("\n")
  end

  defp make_interpreter(model, num_jobs) do
    resolver = TFLiteBEAM.Ops.Builtin.BuiltinResolver.new!()
    builder = InterpreterBuilder.new!(model, resolver)
    interpreter = Interpreter.new!()
    InterpreterBuilder.set_num_threads!(builder, num_jobs)
    :ok = InterpreterBuilder.build!(builder, interpreter)
    Interpreter.set_num_threads!(interpreter, num_jobs)
    interpreter
  end

  defp take_first_and_reshape(tensor) do
    shape = Tuple.delete_at(Nx.shape(tensor), 0)

    Nx.take(tensor, 0)
    |> Nx.reshape(shape)
  end
end