defmodule NxSignal.PeakFinding do
  @moduledoc """
  Peak finding algorithms.
  """

  import Nx.Defn
  import Nx, only: [u8: 1, s64: 1]

  @doc """
  Finds a relative minimum along the selected `:axis`.

  A relative minimum is defined by the element being greater
  than its neighbors along the axis `:axis`.

  Returns a map in the following format:

      %{
        indices: #Nx.Tensor<...>,
        valid_indices: #Nx.Tensor<...>
      }

    * `:indices` - the `{n, rank}` tensor of indices.
      Contains `-1` as a placeholder for invalid indices.

    * `:valid_indices` - the number of valid indices that lead the tensor.

  ## Options

    * `:axis` - the axis along which to do comparisons. Defaults to 0.
    * `:order` - the number of neighbor samples considered for the
      comparison in each direction. Defaults to 1.

  ## Examples

      iex> x = Nx.tensor([2, 1, 2, 3, 2, 0, 1, 0])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmin(x)
      iex> valid_indices
      #Nx.Tensor<
        u32
        2
      >
      iex> indices
      #Nx.Tensor<
        s32[8][1]
        [
          [1],
          [5],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1]
        ]
      >
      iex> Nx.slice_along_axis(indices, 0, Nx.to_number(valid_indices), axis: 0)
      #Nx.Tensor<
        s32[2][1]
        [
          [1],
          [5]
        ]
      >

  For the same tensor in the previous example, we can use `:order` to check if
  the relative maxima are extrema in a wider neighborhood.

      iex> x = Nx.tensor([2, 1, 2, 3, 2, 0, 1, 0])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmin(x, order: 3)
      iex> valid_indices
      #Nx.Tensor<
        u32
        1
      >
      iex> indices
      #Nx.Tensor<
        s32[8][1]
        [
          [1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1]
        ]
      >
      iex> Nx.slice_along_axis(indices, 0, Nx.to_number(valid_indices), axis: 0)
      #Nx.Tensor<
        s32[1][1]
        [
          [1]
        ]
      >

  We can also apply this function to tensors with a larger rank:

      iex> x = Nx.tensor([[1, 2, 1, 2], [6, 2, 0, 0], [5, 3, 4, 4]])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmin(x)
      iex> valid_indices
      #Nx.Tensor<
        u32
        2
      >
      iex> indices[0..1]
      #Nx.Tensor<
        s32[2][2]
        [
          [1, 2],
          [1, 3]
        ]
      >
      iex> %{indices: indices} = NxSignal.PeakFinding.argrelmin(x, axis: 1)
      iex> valid_indices
      #Nx.Tensor<
        u32
        2
      >
      iex> indices[0..1]
      #Nx.Tensor<
        s32[2][2]
        [
          [0, 2],
          [2, 1]
        ]
      >

  """
  @doc type: :peak_finding
  defn argrelmin(data, opts \\ []) do
    opts = keyword!(opts, axis: 0, order: 1)
    argrelextrema(data, &Nx.less/2, opts)
  end

  @doc """
  Finds a relative maximum along the selected `:axis`.

  A relative maximum is defined by the element being greater
  than its neighbors along the axis `:axis`.

  Returns a map in the following format:

      %{
        indices: #Nx.Tensor<...>,
        valid_indices: #Nx.Tensor<...>
      }

    * `:indices` - the `{n, rank}` tensor of indices.
      Contains `-1` as a placeholder for invalid indices.

    * `:valid_indices` - the number of valid indices that lead the tensor.

  ## Options

    * `:axis` - the axis along which to do comparisons. Defaults to 0.
    * `:order` - the number of neighbor samples considered for the
      comparison in each direction. Defaults to 1.

  ## Examples

      iex> x = Nx.tensor([2, 1, 2, 3, 2, 0, 1, 0])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmax(x)
      iex> valid_indices
      #Nx.Tensor<
        u32
        2
      >
      iex> indices
      #Nx.Tensor<
        s32[8][1]
        [
          [3],
          [6],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1]
        ]
      >
      iex> Nx.slice_along_axis(indices, 0, Nx.to_number(valid_indices), axis: 0)
      #Nx.Tensor<
        s32[2][1]
        [
          [3],
          [6]
        ]
      >

  For the same tensor in the previous example, we can use `:order` to check if
  the relative maxima are extrema in a wider neighborhood.

      iex> x = Nx.tensor([2, 1, 2, 3, 2, 0, 1, 0])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmax(x, order: 3)
      iex> valid_indices
      #Nx.Tensor<
        u32
        1
      >
      iex> indices
      #Nx.Tensor<
        s32[8][1]
        [
          [3],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1],
          [-1]
        ]
      >
      iex> Nx.slice_along_axis(indices, 0, Nx.to_number(valid_indices), axis: 0)
      #Nx.Tensor<
        s32[1][1]
        [
          [3]
        ]
      >

  We can also apply this function to tensors with a larger rank:

      iex> x = Nx.tensor([[1, 2, 1, 2], [6, 2, 0, 0], [5, 3, 4, 4]])
      iex> %{indices: indices, valid_indices: valid_indices} = NxSignal.PeakFinding.argrelmax(x)
      iex> valid_indices
      #Nx.Tensor<
        u32
        1
      >
      iex> indices[0]
      #Nx.Tensor<
        s32[2]
        [1, 0]
      >
      iex> %{indices: indices} = NxSignal.PeakFinding.argrelmax(x, axis: 1)
      iex> valid_indices
      #Nx.Tensor<
        u32
        1
      >
      iex> indices[0]
      #Nx.Tensor<
        s32[2]
        [0, 1]
      >

  """
  @doc type: :peak_finding
  defn argrelmax(data, opts \\ []) do
    opts = keyword!(opts, axis: 0, order: 1)
    argrelextrema(data, &Nx.greater/2, opts)
  end

  @doc """
  Finds a relative extrema along the selected `:axis`.

  A relative extremum is defined by the given `comparator_fn`
  function of arity 2 function that returns a boolean tensor.

  This is the function upon which `&argrelmax/2` and `&argrelmin/2`
  are implemented.

  Returns a map in the following format:

      %{
        indices: #Nx.Tensor<...>,
        valid_indices: #Nx.Tensor<...>
      }

    * `:indices` - the `{n, rank}` tensor of indices.
      Contains `-1` as a placeholder for invalid indices.

    * `:valid_indices` - the number of valid indices that lead the tensor.

  ## Options

    * `:axis` - the axis along which to do comparisons. Defaults to 0.
    * `:order` - the number of neighbor samples considered for the
      comparison in each direction. Defaults to 1.

  ## Examples

  First, do read the examples on `argrelmax/2` keeping in mind that
  it is equivalent to `argrelextrema(&1, &Nx.greater/2, &2)`, as well
  as `argrelmin/2` which is equivalent to `argrelextrema(&1, &Nx.less/2, &2)`.

  Having that in mind, we will expand on those concepts by using a custom function.
  For instance, we can change the definition of a relative maximum to one where
  a number is a relative maximum if it is greater than or equal to the double of its
  neighbors, as follows:

      iex> comparator = fn x, y -> Nx.greater_equal(x, Nx.multiply(y, 2)) end
      iex> x = Nx.tensor([0, 1, 3, 2, 0, 1, 0, 0, 0, 2, 1])
      iex> result = NxSignal.PeakFinding.argrelextrema(x, comparator)
      iex> result.valid_indices
      #Nx.Tensor<
        u32
        3
      >
      iex> result.indices[0..2]
      #Nx.Tensor<
        s32[3][1]
        [
          [5],
          [7],
          [9]
        ]
      >

  Same applies for finding local minima. In the next example, we
  find all local minima (i.e. `&Nx.less/2`) that are
  different to the global minimum.

      iex> x = Nx.tensor([0, 1, 0, 2, 1, 3, 0, 1])
      iex> global_minimum = Nx.reduce_min(x)
      iex> comparator = fn x, y ->
      ...> x_not_global = Nx.not_equal(x, global_minimum)
      ...> y_not_global = Nx.not_equal(y, global_minimum)
      ...> both_not_global = Nx.logical_and(x_not_global, y_not_global)
      ...> Nx.logical_and(Nx.less(x, y), both_not_global)
      ...> end
      iex> result = NxSignal.PeakFinding.argrelextrema(x, comparator)
      iex> result.valid_indices
      #Nx.Tensor<
        u32
        1
      >
      iex> result.indices[0..0]
      #Nx.Tensor<
        s32[1][1]
        [
          [4]
        ]
      >
  """
  @doc type: :peak_finding
  defn argrelextrema(data, comparator_fn, opts \\ []) do
    opts = keyword!(opts, axis: 0, order: 1)

    data
    |> boolrelextrema(comparator_fn, opts)
    |> nonzero()
  end

  defnp boolrelextrema(data, comparator_fn, opts \\ []) do
    axis = opts[:axis]
    order = opts[:order]
    locs = Nx.iota({Nx.axis_size(data, axis)})

    ones = Nx.broadcast(u8(1), data.shape)
    [ones, _] = Nx.broadcast_vectors([ones, data])

    {results, _} =
      while {results = ones, {data, locs, halt = u8(0), shift = s64(1)}},
            not halt and shift < order + 1 do
        plus = Nx.take(data, Nx.clip(locs + shift, 0, Nx.size(locs) - 1), axis: axis)
        minus = Nx.take(data, Nx.clip(locs - shift, 0, Nx.size(locs) - 1), axis: axis)
        results = comparator_fn.(data, plus) and results
        results = comparator_fn.(data, minus) and results

        {results, {data, locs, not Nx.any(results), shift + 1}}
      end

    results
  end

  deftransformp nonzero(data) do
    flat_data = Nx.reshape(data, {:auto, 1})

    indices =
      for axis <- 0..(Nx.rank(data) - 1),
          reduce: Nx.broadcast(0, {Nx.axis_size(flat_data, 0), Nx.rank(data)}) do
        %{shape: {n, _}} = indices ->
          iota = data.shape |> Nx.iota(axis: axis) |> Nx.reshape({n, 1})
          Nx.put_slice(indices, [0, axis], iota)
      end

    indices_with_mask =
      Nx.select(
        Nx.broadcast(flat_data, indices.shape),
        indices,
        Nx.broadcast(-1, indices.shape)
      )

    order = Nx.argsort(Nx.squeeze(flat_data, axes: [1]), axis: 0, direction: :desc)

    %{indices: Nx.take(indices_with_mask, order), valid_indices: Nx.sum(flat_data)}
  end
end
