defmodule NxImage do
  @moduledoc """
  Image processing in `Nx`.

  All functions expect images to be tensors in either HWC or CHW order,
  with an arbitrary number of leading batch axes.

  All transformations preserve the input type, rounding if necessary.
  For higher precision, cast the input to floating-point beforehand.
  """

  import Nx.Defn

  @doc """
  Crops an image at the center.

  If the image is too small to be cropped to the desired size, it gets
  padded with zeros.

  ## Options

    * `:channels` - channels location, either `:first` or `:last`.
      Defaults to `:last`

  ## Examples

      iex> image = Nx.iota({4, 4, 1}, type: :u8)
      iex> NxImage.center_crop(image, {2, 2})
      #Nx.Tensor<
        u8[2][2][1]
        [
          [
            [5],
            [6]
          ],
          [
            [9],
            [10]
          ]
        ]
      >

      iex> image = Nx.iota({2, 2, 1}, type: :u8)
      iex> NxImage.center_crop(image, {1, 4})
      #Nx.Tensor<
        u8[1][4][1]
        [
          [
            [0],
            [0],
            [1],
            [0]
          ]
        ]
      >

  """
  @doc type: :transformation
  deftransform center_crop(input, size, opts \\ []) when is_tuple(size) do
    opts = Keyword.validate!(opts, channels: :last)
    validate_image!(input)
    center_crop_n(input, [size: size] ++ opts)
  end

  defnp center_crop_n(input, opts) do
    pad_config =
      transform({input, opts}, fn {input, opts} ->
        for {axis, size, out_size} <- spatial_axes_with_sizes(input, opts),
            reduce: List.duplicate({0, 0, 0}, Nx.rank(input)) do
          pad_config ->
            low = div(size - out_size, 2)
            high = low + out_size
            List.replace_at(pad_config, axis, {-low, high - size, 0})
        end
      end)

    Nx.pad(input, 0, pad_config)
  end

  defnp spatial_axes_with_sizes(input, opts \\ []) do
    {height_axis, width_axis} = spatial_axes(input, channels: opts[:channels])
    {height, width} = size(input, channels: opts[:channels])
    {out_height, out_width} = opts[:size]
    [{height_axis, height, out_height}, {width_axis, width, out_width}]
  end

  # Returns the image size as `{height, width}`.
  defnp size(input, opts \\ []) do
    opts = keyword!(opts, channels: :last)
    {height_axis, width_axis} = spatial_axes(input, channels: opts[:channels])
    {Nx.axis_size(input, height_axis), Nx.axis_size(input, width_axis)}
  end

  @doc """
  Resizes an image.

  ## Options

    * `:method` - the resizing method to use, either of `:nearest`,
      `:bilinear`, `:bicubic`, `:lanczos3`, `:lanczos5`. Defaults to
      `:bilinear`

    * `:channels` - channels location, either `:first` or `:last`.
      Defaults to `:last`

  ## Examples

      iex> image = Nx.iota({2, 2, 1}, type: :u8)
      iex> NxImage.resize(image, {3, 3}, method: :nearest)
      #Nx.Tensor<
        u8[3][3][1]
        [
          [
            [0],
            [1],
            [1]
          ],
          [
            [2],
            [3],
            [3]
          ],
          [
            [2],
            [3],
            [3]
          ]
        ]
      >

      iex> image = Nx.iota({2, 2, 1}, type: :f32)
      iex> NxImage.resize(image, {3, 3}, method: :bilinear)
      #Nx.Tensor<
        f32[3][3][1]
        [
          [
            [0.0],
            [0.5],
            [1.0]
          ],
          [
            [1.0],
            [1.5],
            [2.0]
          ],
          [
            [2.0],
            [2.5],
            [3.0]
          ]
        ]
      >

  """
  @doc type: :transformation
  deftransform resize(input, size, opts \\ []) when is_tuple(size) do
    opts = Keyword.validate!(opts, channels: :last, method: :bilinear)
    validate_image!(input)
    resize_n(input, [size: size] ++ opts)
  end

  defnp resize_n(input, opts) do
    transform({input, opts}, fn {input, opts} ->
      {spatial_axes, out_shape} =
        input
        |> spatial_axes_with_sizes(opts)
        |> Enum.reject(fn {_axis, size, out_size} -> Elixir.Kernel.==(size, out_size) end)
        |> Enum.map_reduce(Nx.shape(input), fn {axis, _size, out_size}, out_shape ->
          {axis, put_elem(out_shape, axis, out_size)}
        end)

      resized_input =
        case opts[:method] do
          :nearest ->
            resize_nearest(input, out_shape, spatial_axes)

          :bilinear ->
            resize_with_kernel(input, out_shape, spatial_axes, &fill_linear_kernel/1)

          :bicubic ->
            resize_with_kernel(input, out_shape, spatial_axes, &fill_cubic_kernel/1)

          :lanczos3 ->
            resize_with_kernel(input, out_shape, spatial_axes, &fill_lanczos_kernel(3, &1))

          :lanczos5 ->
            resize_with_kernel(input, out_shape, spatial_axes, &fill_lanczos_kernel(5, &1))

          method ->
            raise ArgumentError,
                  "expected :method to be either of :nearest, :bilinear, :bicubic, " <>
                    ":lanczos3, :lanczos5, got: #{inspect(method)}"
        end

      cast_to(resized_input, input)
    end)
  end

  defnp spatial_axes(input, opts \\ []) do
    channels = opts[:channels]

    transform({input, channels}, fn {input, channels} ->
      axes =
        case channels do
          :first -> [-2, -1]
          :last -> [-3, -2]
        end

      axes
      |> Enum.map(&Nx.axis_index(input, &1))
      |> List.to_tuple()
    end)
  end

  defnp cast_to(left, right) do
    left_type = Nx.type(left)
    right_type = Nx.type(right)

    left =
      if Nx.Type.float?(left_type) and Nx.Type.integer?(right_type) do
        Nx.round(left)
      else
        left
      end

    left
    |> Nx.as_type(right_type)
    |> Nx.reshape(left, names: Nx.names(right))
  end

  defnp resize_nearest(input, out_shape, spatial_axes) do
    transform({input, out_shape, spatial_axes}, fn {input, out_shape, spatial_axes} ->
      singular_shape = List.duplicate(1, Nx.rank(input)) |> List.to_tuple()

      for axis <- spatial_axes, reduce: input do
        input ->
          input_shape = Nx.shape(input)
          input_size = elem(input_shape, axis)
          output_size = elem(out_shape, axis)
          inv_scale = input_size / output_size
          offset = (Nx.iota({output_size}) + 0.5) * inv_scale
          offset = offset |> Nx.floor() |> Nx.as_type({:s, 32})

          offset =
            offset
            |> Nx.reshape(put_elem(singular_shape, axis, output_size))
            |> Nx.broadcast(put_elem(input_shape, axis, output_size))

          Nx.take_along_axis(input, offset, axis: axis)
      end
    end)
  end

  @f32_eps :math.pow(2, -23)

  defnp resize_with_kernel(input, out_shape, spatial_axes, kernel_fun) do
    transform({input, out_shape, spatial_axes}, fn {input, out_shape, spatial_axes} ->
      for axis <- spatial_axes, reduce: input do
        input ->
          input_shape = Nx.shape(input)
          input_size = elem(input_shape, axis)
          output_size = elem(out_shape, axis)

          inv_scale = input_size / output_size
          kernel_scale = Nx.max(1, inv_scale)

          sample_f = (Nx.iota({1, output_size}) + 0.5) * inv_scale - 0.5
          x = Nx.abs(sample_f - Nx.iota({input_size, 1})) / kernel_scale
          weights = kernel_fun.(x)

          weights_sum = Nx.sum(weights, axes: [0], keep_axes: true)

          weights =
            Nx.select(Nx.abs(weights) > 1000 * @f32_eps, safe_divide(weights, weights_sum), 0)

          input = Nx.dot(input, [axis], weights, [0])
          # The transformed axis is moved to the end, so we transpose back
          reorder_axis(input, -1, axis)
      end
    end)
  end

  defnp fill_linear_kernel(x) do
    Nx.max(0, 1 - x)
  end

  defnp fill_cubic_kernel(x) do
    # See https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm
    out = (1.5 * x - 2.5) * x * x + 1
    out = Nx.select(x >= 1, ((-0.5 * x + 2.5) * x - 4) * x + 2, out)
    Nx.select(x >= 2, 0, out)
  end

  @pi :math.pi()

  defnp fill_lanczos_kernel(radius, x) do
    y = radius * Nx.sin(@pi * x) * Nx.sin(@pi * x / radius)
    out = Nx.select(x > 1.0e-3, safe_divide(y, @pi ** 2 * x ** 2), 1)
    Nx.select(x > radius, 0, out)
  end

  defnp safe_divide(x, y) do
    x / Nx.select(y != 0, y, 1)
  end

  defnp reorder_axis(tensor, axis, target_axis) do
    transform({tensor, axis, target_axis}, fn {tensor, axis, target_axis} ->
      axes = Nx.axes(tensor)
      {source_axis, axes} = List.pop_at(axes, axis)
      axes = List.insert_at(axes, target_axis, source_axis)
      Nx.transpose(tensor, axes: axes)
    end)
  end

  @doc """
  Scales an image such that the short edge matches the given size.

  ## Options

    * `:method` - the resizing method to use, same as `resize/2`

    * `:channels` - channels location, either `:first` or `:last`.
      Defaults to `:last`

  ## Examples

      iex> image = Nx.iota({2, 4, 1}, type: :u8)
      iex> resized_image = NxImage.resize_short(image, 3, method: :nearest)
      iex> Nx.shape(resized_image)
      {3, 6, 1}

      iex> image = Nx.iota({4, 2, 1}, type: :u8)
      iex> resized_image = NxImage.resize_short(image, 3, method: :nearest)
      iex> Nx.shape(resized_image)
      {6, 3, 1}

  """
  @doc type: :transformation
  deftransform resize_short(input, size, opts \\ []) when is_integer(size) do
    opts = Keyword.validate!(opts, channels: :last, method: :bilinear)
    validate_image!(input)
    resize_short_n(input, [size: size] ++ opts)
  end

  defnp resize_short_n(input, opts) do
    size = opts[:size]
    method = opts[:method]
    channels = opts[:channels]

    {height, width} = size(input, channels: channels)
    {out_height, out_width} = transform({height, width, size}, &resize_short_size/1)

    resize(input, {out_height, out_width}, method: method, channels: channels)
  end

  defp resize_short_size({height, width, size}) do
    {short, long} = if height < width, do: {height, width}, else: {width, height}

    out_short = size
    out_long = floor(size * long / short)

    if height < width, do: {out_short, out_long}, else: {out_long, out_short}
  end

  @doc """
  Normalizes an image according to the given per-channel mean and
  standard deviation.

    * `:channels` - channels location, either `:first` or `:last`.
      Defaults to `:last`

  ## Examples

      iex> image = Nx.iota({2, 2, 3}, type: :f32)
      iex> mean = Nx.tensor([0.485, 0.456, 0.406])
      iex> std = Nx.tensor([0.229, 0.224, 0.225])
      iex> NxImage.normalize(image, mean, std)
      #Nx.Tensor<
        f32[2][2][3]
        [
          [
            [-2.1179039478302, 2.4285714626312256, 7.084444522857666],
            [10.982532501220703, 15.821427345275879, 20.41777801513672]
          ],
          [
            [24.08296775817871, 29.214284896850586, 33.7511100769043],
            [37.183406829833984, 42.607139587402344, 47.08444595336914]
          ]
        ]
      >

  """
  @doc type: :transformation
  defn normalize(input, mean, std, opts \\ []) do
    opts = keyword!(opts, channels: :last)
    validate_image!(input)

    mean = broadcast_channel_info(mean, input, opts[:channels], "mean")
    std = broadcast_channel_info(std, input, opts[:channels], "std")

    normalized_input = (input - mean) / std

    cast_to(normalized_input, input)
  end

  deftransformp broadcast_channel_info(tensor, input, channels, name) do
    rank = Nx.rank(input)

    channels_axis =
      case channels do
        :first -> rank - 3
        :last -> rank - 1
      end

    num_channels = Nx.axis_size(input, channels_axis)

    case Nx.shape(tensor) do
      {^num_channels} ->
        :ok

      shape ->
        raise ArgumentError,
              "expected #{name} to have shape {#{num_channels}}, got: #{inspect(shape)}"
    end

    shape = 1 |> Tuple.duplicate(rank) |> put_elem(channels_axis, :auto)
    Nx.reshape(tensor, shape)
  end

  @doc """
  Converts pixel values (0-255) into a continuous range.

  ## Examples

      iex> image = Nx.tensor([[[0], [128]], [[191], [255]]])
      iex> NxImage.to_continuous(image, 0.0, 1.0)
      #Nx.Tensor<
        f32[2][2][1]
        [
          [
            [0.0],
            [0.501960813999176]
          ],
          [
            [0.7490196228027344],
            [1.0]
          ]
        ]
      >

      iex> image = Nx.tensor([[[0], [128]], [[191], [255]]])
      iex> NxImage.to_continuous(image, -1.0, 1.0)
      #Nx.Tensor<
        f32[2][2][1]
        [
          [
            [-1.0],
            [0.003921627998352051]
          ],
          [
            [0.49803924560546875],
            [1.0]
          ]
        ]
      >

  """
  @doc type: :conversion
  defn to_continuous(input, min, max) do
    validate_image!(input)

    input / 255.0 * (max - min) + min
  end

  @doc """
  Converts values from continuous range into pixel values (0-255).

  ## Examples

      iex> image = Nx.tensor([[[0.0], [0.5]], [[0.75], [1.0]]])
      iex> NxImage.from_continuous(image, 0.0, 1.0)
      #Nx.Tensor<
        u8[2][2][1]
        [
          [
            [0],
            [128]
          ],
          [
            [191],
            [255]
          ]
        ]
      >

      iex> image = Nx.tensor([[[-1.0], [0.0]], [[0.5], [1.0]]])
      iex> NxImage.from_continuous(image, -1.0, 1.0)
      #Nx.Tensor<
        u8[2][2][1]
        [
          [
            [0],
            [128]
          ],
          [
            [191],
            [255]
          ]
        ]
      >

  """
  @doc type: :conversion
  defn from_continuous(input, min, max) do
    validate_image!(input)

    input = (input - min) / (max - min) * 255.0

    input
    |> Nx.round()
    |> Nx.clip(0, 255)
    |> Nx.as_type(:u8)
  end

  deftransformp validate_image!(input) do
    rank = Nx.rank(input)

    if rank < 3 do
      raise ArgumentError,
            "expected the image input to have rank 3 or higher, got: #{inspect(rank)}"
    end
  end
end