defmodule TfIdfTest do
  use ExUnit.Case
  doctest TfIdf

  test "generates word rankings map" do
    assert """
           This is a test document.
           This is another test document.
           This is a third test document.
           """
           |> String.split("\n", trim: true)
           |> TfIdf.build_model() == %TfIdf.Model{
             count: 3,
             docs: [
               {"This is a test document.",
                %TfIdf.Doc{
                  tf: %{"" => 1, "a" => 1, "document" => 1, "is" => 1, "test" => 1, "this" => 1},
                  count: 6
                }},
               {"This is another test document.",
                %TfIdf.Doc{
                  tf: %{
                    "" => 1,
                    "another" => 1,
                    "document" => 1,
                    "is" => 1,
                    "test" => 1,
                    "this" => 1
                  },
                  count: 6
                }},
               {"This is a third test document.",
                %TfIdf.Doc{
                  tf: %{
                    "" => 1,
                    "a" => 1,
                    "document" => 1,
                    "is" => 1,
                    "test" => 1,
                    "third" => 1,
                    "this" => 1
                  },
                  count: 7
                }}
             ],
             idf: %{
               "" => 0.0,
               "a" => 0.12493873660829992,
               "another" => 0.3010299956639812,
               "document" => 0.0,
               "is" => 0.0,
               "test" => 0.0,
               "third" => 0.3010299956639812,
               "this" => 0.0
             }
           }
  end

  test "wikipedia example tf_idf" do
    model =
      """
      This is a a sample
      This is another another example example example
      """
      |> String.split("\n", trim: true)
      |> TfIdf.build_model()

    assert TfIdf.calculate_tf_idf(model, "example") ==
             [{"This is another another example example example", 0.0754676824524348}]

    assert TfIdf.calculate_tf_idf(model, "sample") ==
             [{"This is a a sample", 0.03521825181113625}]

    assert TfIdf.calculate_tf_idf(model, "this") == []
  end

  test "wikipedia example dn_idf" do
    model =
      """
      This is a a sample
      This is another another example example example
      """
      |> String.split("\n", trim: true)
      |> TfIdf.build_model()

    assert TfIdf.calculate_dn_idf(model, "example") ==
             [
               {"This is another another example example example", 0.17609125905568124},
               {"This is a a sample", 0.08804562952784062}
             ]

    assert TfIdf.calculate_dn_idf(model, "sample") ==
             [
               {"This is a a sample", 0.13206844429176093},
               {"This is another another example example example", 0.08804562952784062}
             ]

    assert(TfIdf.calculate_dn_idf(model, "this") == [])
  end

  test "wikipedia example log_norm_idf" do
    model =
      """
      This is a a sample
      This is another another example example example
      """
      |> String.split("\n", trim: true)
      |> TfIdf.build_model()

    assert TfIdf.calculate_log_norm_idf(model, "example") ==
             [
               {"This is another another example example example", 0.11129376362989402}
             ]

    assert TfIdf.calculate_log_norm_idf(model, "sample") ==
             [
               {"This is a a sample", 0.05300875094999672}
             ]

    assert(TfIdf.calculate_log_norm_idf(model, "this") == [])
  end
end