extorch_vision

defmodule ExTorch.Vision do
  @moduledoc """
  TorchVision ops for ExTorch.

  Loads `libtorchvision.so` on first use and registers its custom ops
  (roi_align, nms, deform_conv2d, etc.) with `ExTorch.Export.OpRegistry`
  so they are available when running exported models that use torchvision
  operators.

  `libtorchvision.so` is built from source automatically during
  `mix compile` -- no manual steps needed.

  ## Setup

  Call `ExTorch.Vision.setup!/0` once, or let it auto-initialize on first
  op call. The library is resolved in this order:

  1. Application config: `config :extorch_vision, library_path: "..."`
  2. Environment variable: `TORCHVISION_LIB_PATH=/path/to/libtorchvision.so`
  3. Auto-built library in `priv/` (from `mix compile`)

  ## Usage

      ExTorch.Vision.nms(boxes, scores, 0.5)
      ExTorch.Vision.roi_align(input, rois, 1.0, 7, 7)
      ExTorch.Vision.deform_conv2d(input, weight, offset, mask, bias, 1, 1, 0, 0)

  ## Local development

  To develop against a local ExTorch checkout instead of the Hex release:

      export EXTORCH_PATH=../extorch
      mix deps.get
      mix test
  """

  @setup_key {__MODULE__, :initialized}

  @doc """
  Load the torchvision shared library and register ops with the
  `ExTorch.Export.OpRegistry`.

  Safe to call multiple times — subsequent calls are no-ops.
  Also called automatically on first op call if not yet initialized.
  """
  def setup! do
    case :persistent_term.get(@setup_key, false) do
      true ->
        :ok

      false ->
        path = library_path()
        # Pre-load runtime dependencies that libtorchvision.so needs but
        # may not have as DT_NEEDED entries (e.g., nvjpeg when built with
        # NVJPEG_FOUND but the standalone cmake doesn't add the link dep).
        preload_dependencies()
        ExTorch.Native.load_torch_library(path)
        ExTorch.Export.OpRegistry.register(ExTorch.Vision.Ops)
        :persistent_term.put(@setup_key, true)
        :ok
    end
  end

  @doc """
  Resolve the path to the torchvision shared library.

  Checks (in order):
  1. `:extorch_vision` application config `:library_path`
  2. `TORCHVISION_LIB_PATH` environment variable
  3. Built-from-source `libtorchvision.so` in priv/
  """
  def library_path do
    Application.get_env(:extorch_vision, :library_path) ||
      System.get_env("TORCHVISION_LIB_PATH") ||
      built_lib_path()
  end

  @doc """
  Non-maximum suppression.

  Filters overlapping bounding boxes by score, removing boxes with
  IoU > `iou_threshold` against higher-scoring boxes.

  ## Parameters
    * `boxes` -- `{N, 4}` tensor of `[x1, y1, x2, y2]` boxes
    * `scores` -- `{N}` tensor of scores
    * `iou_threshold` -- IoU threshold (float)

  Returns a 1-D tensor of indices to keep.
  """
  def nms(boxes, scores, iou_threshold) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::nms", "", [
      {:tensor, boxes},
      {:tensor, scores},
      {:float, iou_threshold / 1}
    ])
  end

  @doc """
  Region of Interest Align.

  Extracts fixed-size feature maps from regions of interest in a feature map
  using bilinear interpolation.

  ## Parameters
    * `input` -- `{N, C, H, W}` feature map
    * `rois` -- `{K, 5}` tensor of `[batch_index, x1, y1, x2, y2]`
    * `spatial_scale` -- scale factor mapping input coords to feature map coords
    * `output_h` -- output height
    * `output_w` -- output width
    * `sampling_ratio` -- number of sampling points (default: -1 for adaptive)
    * `aligned` -- use half-pixel offset (default: false)
  """
  def roi_align(
        input,
        rois,
        spatial_scale,
        output_h,
        output_w,
        sampling_ratio \\ -1,
        aligned \\ false
      ) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::roi_align", "", [
      {:tensor, input},
      {:tensor, rois},
      {:float, spatial_scale / 1},
      {:int, output_h},
      {:int, output_w},
      {:int, sampling_ratio},
      {:bool, aligned}
    ])
  end

  @doc """
  Region of Interest Pooling.

  Extracts fixed-size feature maps from regions of interest using max pooling.
  Returns `{output, argmax}` where argmax records which input element was
  selected for each output element.

  ## Parameters
    * `input` -- `{N, C, H, W}` feature map
    * `rois` -- `{K, 5}` tensor of `[batch_index, x1, y1, x2, y2]`
    * `spatial_scale` -- scale factor mapping input coords to feature map coords
    * `output_h` -- output height
    * `output_w` -- output width
  """
  def roi_pool(input, rois, spatial_scale, output_h, output_w) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::roi_pool", "", [
      {:tensor, input},
      {:tensor, rois},
      {:float, spatial_scale / 1},
      {:int, output_h},
      {:int, output_w}
    ])
  end

  @doc """
  Position-Sensitive Region of Interest Align.

  A variant of ROI Align for position-sensitive score maps, used in R-FCN
  architectures. Returns `{output, channel_mapping}`.

  ## Parameters
    * `input` -- `{N, C, H, W}` feature map
    * `rois` -- `{K, 5}` tensor of `[batch_index, x1, y1, x2, y2]`
    * `spatial_scale` -- scale factor mapping input coords to feature map coords
    * `output_h` -- output height
    * `output_w` -- output width
    * `sampling_ratio` -- number of sampling points (default: -1 for adaptive)
  """
  def ps_roi_align(input, rois, spatial_scale, output_h, output_w, sampling_ratio \\ -1) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::ps_roi_align", "", [
      {:tensor, input},
      {:tensor, rois},
      {:float, spatial_scale / 1},
      {:int, output_h},
      {:int, output_w},
      {:int, sampling_ratio}
    ])
  end

  @doc """
  Position-Sensitive Region of Interest Pooling.

  A variant of ROI Pooling for position-sensitive score maps.
  Returns `{output, channel_mapping}`.

  ## Parameters
    * `input` -- `{N, C, H, W}` feature map
    * `rois` -- `{K, 5}` tensor of `[batch_index, x1, y1, x2, y2]`
    * `spatial_scale` -- scale factor mapping input coords to feature map coords
    * `output_h` -- output height
    * `output_w` -- output width
  """
  def ps_roi_pool(input, rois, spatial_scale, output_h, output_w) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::ps_roi_pool", "", [
      {:tensor, input},
      {:tensor, rois},
      {:float, spatial_scale / 1},
      {:int, output_h},
      {:int, output_w}
    ])
  end

  @doc """
  Deformable Convolution v2.

  Performs a 2D convolution with learnable offsets and optional modulation
  masks, enabling the network to learn spatially-varying receptive fields.

  ## Parameters
    * `input` -- `{N, C_in, H, W}` input tensor
    * `weight` -- `{C_out, C_in/groups, kH, kW}` convolution weights
    * `offset` -- `{N, 2*kH*kW*offset_groups, H_out, W_out}` offset field
    * `mask` -- `{N, kH*kW*offset_groups, H_out, W_out}` modulation mask,
      or pass a zero tensor with `use_mask: false`
    * `bias` -- `{C_out}` bias tensor, or zero tensor for no bias
    * `stride_h`, `stride_w` -- convolution stride
    * `pad_h`, `pad_w` -- padding
    * `dilation_h`, `dilation_w` -- dilation
    * `groups` -- number of groups for grouped convolution
    * `offset_groups` -- number of groups for offset computation
    * `use_mask` -- whether to apply the modulation mask (default: true)
  """
  def deform_conv2d(
        input,
        weight,
        offset,
        mask,
        bias,
        stride_h,
        stride_w,
        pad_h,
        pad_w,
        dilation_h \\ 1,
        dilation_w \\ 1,
        groups \\ 1,
        offset_groups \\ 1,
        use_mask \\ true
      ) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("torchvision::deform_conv2d", "", [
      {:tensor, input},
      {:tensor, weight},
      {:tensor, offset},
      {:tensor, mask},
      {:tensor, bias},
      {:int, stride_h},
      {:int, stride_w},
      {:int, pad_h},
      {:int, pad_w},
      {:int, dilation_h},
      {:int, dilation_w},
      {:int, groups},
      {:int, offset_groups},
      {:bool, use_mask}
    ])
  end

  # ===========================================================================
  # Image I/O ops (image:: namespace)
  # ===========================================================================

  @doc """
  Read a file from disk into a uint8 tensor.

  Returns a 1-D `{N}` tensor of bytes.
  """
  def read_file(path) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::read_file", "", [
      {:tensor, ExTorch.Native.from_binary(path, {byte_size(path)}, :uint8)}
    ])
  end

  @doc """
  Write a uint8 tensor to a file on disk.
  """
  def write_file(path, data) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::write_file", "", [
      {:tensor, ExTorch.Native.from_binary(path, {byte_size(path)}, :uint8)},
      {:tensor, data}
    ])
  end

  @doc """
  Decode a JPEG image from an in-memory uint8 tensor.

  ## Parameters
    * `data` -- 1-D uint8 tensor of JPEG bytes
    * `mode` -- color mode: 0 = unchanged, 1 = grayscale, 2 = RGB (default: 0)
    * `apply_exif` -- apply EXIF orientation (default: false)

  Returns a `{C, H, W}` uint8 tensor.
  """
  def decode_jpeg(data, mode \\ 0, apply_exif \\ false) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::decode_jpeg", "", [
      {:tensor, data},
      {:int, mode},
      {:bool, apply_exif}
    ])
  end

  @doc """
  Encode a `{C, H, W}` uint8 tensor as JPEG bytes.

  Returns a 1-D uint8 tensor of JPEG data.
  """
  def encode_jpeg(image, quality \\ 75) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::encode_jpeg", "", [
      {:tensor, image},
      {:int, quality}
    ])
  end

  @doc """
  Decode a PNG image from an in-memory uint8 tensor.

  ## Parameters
    * `data` -- 1-D uint8 tensor of PNG bytes
    * `mode` -- color mode: 0 = unchanged, 1 = grayscale, 2 = RGB (default: 0)
    * `apply_exif` -- apply EXIF orientation (default: false)

  Returns a `{C, H, W}` uint8 tensor.
  """
  def decode_png(data, mode \\ 0, apply_exif \\ false) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::decode_png", "", [
      {:tensor, data},
      {:int, mode},
      {:bool, apply_exif}
    ])
  end

  @doc """
  Encode a `{C, H, W}` uint8 tensor as PNG bytes.

  Returns a 1-D uint8 tensor of PNG data.
  """
  def encode_png(image, compression \\ 6) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::encode_png", "", [
      {:tensor, image},
      {:int, compression}
    ])
  end

  @doc """
  Decode a WebP image from an in-memory uint8 tensor.

  Returns a `{C, H, W}` uint8 tensor.
  """
  def decode_webp(data, mode \\ 0) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::decode_webp", "", [
      {:tensor, data},
      {:int, mode}
    ])
  end

  @doc """
  Decode a GIF image from an in-memory uint8 tensor.

  Returns a `{T, C, H, W}` uint8 tensor for animated GIFs,
  or `{C, H, W}` for single-frame.
  """
  def decode_gif(data) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::decode_gif", "", [
      {:tensor, data}
    ])
  end

  @doc """
  Decode any supported image format from an in-memory uint8 tensor.

  Automatically detects the format (JPEG, PNG, GIF, WebP).
  Returns a `{C, H, W}` uint8 tensor.
  """
  def decode_image(data, mode \\ 0, apply_exif \\ false) do
    ensure_setup!()

    ExTorch.Native.dispatch_op("image::decode_image", "", [
      {:tensor, data},
      {:int, mode},
      {:bool, apply_exif}
    ])
  end

  @doc """
  Batch-decode JPEG images on GPU using NVJPEG.

  Uses the NVJPEG library for GPU-accelerated JPEG decoding. Available
  on all CUDA GPUs. A100+ GPUs have a dedicated hardware JPEG decoder
  for additional speedup.

  ## Parameters
    * `encoded_images` -- list of 1-D uint8 tensors, each containing JPEG bytes
    * `mode` -- color mode: 0 = unchanged, 1 = grayscale, 2 = RGB (default: 0)
    * `device` -- CUDA device (default: `:cuda`)

  Returns a list of `{C, H, W}` uint8 tensors on the specified device.
  """
  def decode_jpegs_cuda(encoded_images, mode \\ 0, device \\ :cuda) do
    ensure_setup!()
    tensor_args = Enum.map(encoded_images, &{:tensor, &1})

    ExTorch.Native.dispatch_op("image::decode_jpegs_cuda", "", [
      {:list, tensor_args},
      {:int, mode},
      {:device, device_to_arg(device)}
    ])
  end

  @doc """
  Batch-encode images as JPEG on GPU using hardware acceleration.

  ## Parameters
    * `images` -- list of `{C, H, W}` uint8 CUDA tensors
    * `quality` -- JPEG quality 1-100 (default: 75)

  Returns a list of 1-D uint8 tensors containing JPEG bytes.
  """
  def encode_jpegs_cuda(images, quality \\ 75) do
    ensure_setup!()
    tensor_args = Enum.map(images, &{:tensor, &1})

    ExTorch.Native.dispatch_op("image::encode_jpegs_cuda", "", [
      {:list, tensor_args},
      {:int, quality}
    ])
  end

  defp device_to_arg(:cuda), do: {"cuda", 0}
  defp device_to_arg({:cuda, idx}), do: {"cuda", idx}
  defp device_to_arg(:cpu), do: "cpu"

  defp ensure_setup!, do: setup!()

  # Pre-load shared libraries that libtorchvision.so depends on at runtime.
  # The standalone cmake build may compile with NVJPEG_FOUND but not add
  # libnvjpeg.so as a DT_NEEDED entry. Loading it first via dlopen with
  # RTLD_GLOBAL makes the symbols available for libtorchvision.
  defp preload_dependencies do
    cuda_home = ExTorch.Vision.Build.cuda_home()

    if cuda_home do
      nvjpeg = Path.join([cuda_home, "lib64", "libnvjpeg.so"])

      if File.exists?(nvjpeg) do
        try do
          ExTorch.Native.load_torch_library(nvjpeg)
        rescue
          # Not critical — nvjpeg ops will fail gracefully
          _ -> :ok
        end
      end
    end
  end

  defp built_lib_path do
    path = ExTorch.Vision.Build.lib_path()

    if File.exists?(path) do
      path
    else
      raise """
      libtorchvision.so not found. Run `mix torchvision.build` first, or set one of:

        - config :extorch_vision, library_path: "/path/to/libtorchvision.so"
        - TORCHVISION_LIB_PATH=/path/to/libtorchvision.so
      """
    end
  end
end