tf_idf

defmodule TfIdf do
  alias :math, as: Math

  defmodule Doc do
    defstruct tf: %{}, max_tf: 1, count: 0
  end

  defmodule Model do
    defstruct docs: [], count: 0, idf: %{}
  end

  def build_model(docs) do
    docs
    |> calculate_docs_tf()
    |> calculate_idf()
  end

  def calculate_tf_idf(model, query) do
    model.docs
    |> Enum.reduce([], fn {doc_name, %{tf: tf, count: total_count}}, acc ->
      score =
        query
        |> String.downcase()
        |> String.split(~r/[^a-zA-Z]/)
        |> Enum.map(fn token ->
          idf = Map.get(model.idf, token, 0)
          token_count = Map.get(tf, token, 0)
          tf = token_count / total_count
          tf * idf
        end)
        |> Enum.sum()

      if score == 0 do
        acc
      else
        [{doc_name, score} | acc]
      end
    end)
    |> Enum.sort_by(fn {_, score} -> score end, :desc)
  end

  def calculate_log_norm_idf(model, query) do
    model.docs
    |> Enum.reduce([], fn {doc_name, %{tf: tf, count: total_count}}, acc ->
      score =
        query
        |> String.downcase()
        |> String.split(~r/[^a-zA-Z]/)
        |> Enum.map(fn token ->
          idf = Map.get(model.idf, token, 0)
          token_count = Map.get(tf, token, 0)
          tf = token_count / total_count

          if tf == 0 do
            0
          else
            (1 + Math.log10(tf)) * idf
          end
        end)
        |> Enum.sum()

      if score == 0 do
        acc
      else
        [{doc_name, score} | acc]
      end
    end)
    |> Enum.sort_by(fn {_, score} -> score end, :desc)
  end

  def calculate_dn_idf(model, query) do
    model.docs
    |> Enum.reduce([], fn {doc_name, %{tf: tf, max_tf: max_tf}}, acc ->
      score =
        query
        |> String.downcase()
        |> String.split(~r/[^a-zA-Z]/)
        |> Enum.map(fn token ->
          idf = Map.get(model.idf, token, 0)
          token_count = Map.get(tf, token, 0)
          tf = 0.5 + 0.5 * (token_count / max_tf)
          tf * idf
        end)
        |> Enum.sum()

      if score == 0 do
        acc
      else
        [{doc_name, score} | acc]
      end
    end)
    |> Enum.sort_by(fn {_, score} -> score end, :desc)
  end

  def calculate_token_value(model, token) do
    idf = Map.get(model.idf, token, 0)

    model.docs
    |> Enum.map(fn {doc_name, %{tf: tf, count: total_count}} ->
      token_count = Map.get(tf, token, 0)
      tf = token_count / total_count
      {doc_name, tf * idf}
    end)
  end

  defp calculate_docs_tf(docs) do
    Enum.map(docs, &calculate_doc_tf/1)
  end

  defp calculate_doc_tf({name, doc}) when is_binary(name) and is_binary(doc) do
    {_, tf} = calculate_doc_tf(doc)
    {name, tf}
  end

  defp calculate_doc_tf(doc) when is_binary(doc) do
    words =
      doc
      |> String.downcase()
      |> String.split(~r/[^a-zA-Z]/)

    n = length(words)
    tf = Enum.frequencies(words)
    max_tf = Map.values(tf) |> Enum.max()

    {doc, %TfIdf.Doc{tf: tf, max_tf: max_tf, count: n}}
  end

  defp calculate_idf(docs) do
    n = length(docs)

    idf =
      docs
      |> Enum.flat_map(&Map.keys(elem(&1, 1).tf))
      |> Enum.frequencies()
      |> Map.new(fn {t, c} -> {t, Math.log10((n + 1) / (c + 1))} end)

    %TfIdf.Model{docs: docs, count: n, idf: idf}
  end
end