ex_maxsim_cpu

defmodule ExMaxsimCpu do
  @moduledoc """
  High-performance MaxSim (Maximum Similarity) scoring using BLAS GEMM operations.

  MaxSim computes the maximum similarity score between a query and a set of documents.
  For each query token, it finds the maximum similarity with any document token,
  then sums these maximums to produce the final score.

  ## Usage with Nx Tensors

  ExMaxsimCpu exposes an Nx-only public API.

      # Query: [q_len, dim] tensor
      query = Nx.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], type: :f32)

      # Documents: [n_docs, d_len, dim] tensor
      docs = Nx.tensor([
        [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
        [[0.7, 0.8, 0.9], [0.1, 0.0, 0.1]]
      ], type: :f32)

      scores = ExMaxsimCpu.maxsim_scores(query, docs)
      # => #Nx.Tensor<f32[2]>

  ## Performance Notes

  - Uses Dirty CPU schedulers for compute-intensive operations
  - Leverages BLAS (OpenBLAS on Linux, Accelerate on macOS) for matrix operations
  - SIMD-optimized max reduction (AVX2 on x86_64, NEON on ARM64)
  - Rayon for parallel document processing

  ## Environment Variables

  - `RAYON_NUM_THREADS`: Control Rayon parallelism (default: number of CPUs)
  - `OPENBLAS_NUM_THREADS`: Set to 1 to avoid oversubscription with Rayon

  ## Input Requirements

  - Embeddings must be L2-normalized per token for correct MaxSim scores
  """

  alias ExMaxsimCpu.Nif

  @doc """
  Compute MaxSim scores between a query and a batch of documents.

  ## Parameters

  - `query`: An Nx tensor of shape `{q_len, dim}` with type `:f32`
  - `docs`: An Nx tensor of shape `{n_docs, d_len, dim}` with type `:f32`

  ## Returns

  An Nx tensor of shape `{n_docs}` containing the MaxSim score for each document.

  ## Examples

      iex> query = Nx.tensor([[1.0, 0.0], [0.0, 1.0]], type: :f32)
      iex> docs = Nx.tensor([[[1.0, 0.0], [0.0, 1.0]]], type: :f32)
      iex> ExMaxsimCpu.maxsim_scores(query, docs)
      #Nx.Tensor<
        f32[1]
        [2.0]
      >
  """
  @spec maxsim_scores(Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t()
  def maxsim_scores(query, docs) do
    # Validate shapes
    {q_len, q_dim} = validate_query_shape(query)
    {n_docs, d_len, d_dim} = validate_docs_shape(docs)

    if q_dim != d_dim do
      raise ArgumentError,
            "Dimension mismatch: query dim #{q_dim} vs docs dim #{d_dim}"
    end

    # Ensure f32 type and contiguous layout
    query = Nx.as_type(query, :f32)
    docs = Nx.as_type(docs, :f32)

    query_bin = Nx.to_binary(query)
    docs_bin = Nx.to_binary(docs)

    scores_bin = Nif.maxsim_scores_nif(query_bin, q_len, q_dim, docs_bin, n_docs, d_len)

    Nx.from_binary(scores_bin, :f32)
  end

  @doc """
  Compute MaxSim scores for variable-length documents.

  ## Parameters

  - `query`: An Nx tensor of shape `{q_len, dim}` with type `:f32`
  - `docs`: A list of Nx tensors, each with shape `{doc_len_i, dim}` with type `:f32`

  ## Returns

  An Nx tensor of shape `{length(docs)}` containing the MaxSim score for each document.
  """
  @spec maxsim_scores_variable(Nx.Tensor.t(), [Nx.Tensor.t()]) :: Nx.Tensor.t()
  def maxsim_scores_variable(_query, []) do
    raise ArgumentError,
          "Empty document list not supported (Nx cannot create empty tensors)."
  end

  def maxsim_scores_variable(query, docs) when is_list(docs) do
    {q_len, q_dim} = validate_query_shape(query)

    # Validate and convert all documents
    {doc_bins, doc_lens} =
      docs
      |> Enum.with_index()
      |> Enum.map(fn {doc, idx} ->
        shape = Nx.shape(doc)

        case shape do
          {doc_len, dim} when dim == q_dim ->
            doc = Nx.as_type(doc, :f32)
            {Nx.to_binary(doc), doc_len}

          {_, dim} ->
            raise ArgumentError,
                  "Dimension mismatch at doc #{idx}: query dim #{q_dim} vs doc dim #{dim}"

          other ->
            raise ArgumentError,
                  "Invalid doc shape at #{idx}: expected {doc_len, dim}, got #{inspect(other)}"
        end
      end)
      |> Enum.unzip()

    query = Nx.as_type(query, :f32)
    query_bin = Nx.to_binary(query)

    scores_bin = Nif.maxsim_scores_variable_nif(query_bin, q_len, q_dim, doc_bins, doc_lens)

    Nx.from_binary(scores_bin, :f32)
  end


  # Private helpers

  defp validate_query_shape(query) do
    case Nx.shape(query) do
      {q_len, dim} ->
        {q_len, dim}

      other ->
        raise ArgumentError,
              "Invalid query shape: expected {q_len, dim}, got #{inspect(other)}"
    end
  end

  defp validate_docs_shape(docs) do
    case Nx.shape(docs) do
      {n_docs, d_len, dim} ->
        {n_docs, d_len, dim}

      other ->
        raise ArgumentError,
              "Invalid docs shape: expected {n_docs, d_len, dim}, got #{inspect(other)}"
    end
  end
end