defmodule Rsh26pool.Voting.HungarianTest do
  use ExUnit.Case, async: true

  alias Rsh26pool.Voting.Hungarian

  describe "min_cost_assignment/1" do
    test "empty matrix yields no cost and no assignment" do
      assert Hungarian.min_cost_assignment([]) == {0, []}
    end

    test "raises when there are more rows than columns" do
      assert_raise ArgumentError, fn ->
        Hungarian.min_cost_assignment([[1, 2], [3, 4], [5, 6]])
      end
    end

    test "square matrix matches the brute-force optimum" do
      cost = [[4, 1, 3], [2, 0, 5], [3, 2, 2]]
      {total, assignment} = Hungarian.min_cost_assignment(cost)

      {brute_total, _} = brute_force(cost)
      assert total == brute_total
      assert valid_assignment?(assignment, 3, 3)
      assert picked_sum(cost, assignment) == total
    end

    test "rectangular matrix (rows < cols) matches the brute-force optimum" do
      cost = [[3, 1, 2, 4], [4, 2, 1, 0]]
      {total, assignment} = Hungarian.min_cost_assignment(cost)

      {brute_total, _} = brute_force(cost)
      assert total == brute_total
      assert valid_assignment?(assignment, 2, 4)
      assert picked_sum(cost, assignment) == total
    end

    test "matches the brute-force optimum across many random small matrices" do
      # Deterministic: fixed seed, and the failing matrix is printed for repro.
      :rand.seed(:exsss, {13, 57, 911})

      for _ <- 1..250 do
        n = Enum.random(1..4)
        m = Enum.random(n..5)
        cost = for _ <- 1..n, do: for(_ <- 1..m, do: Enum.random(0..9))

        {total, assignment} = Hungarian.min_cost_assignment(cost)
        {brute_total, _} = brute_force(cost)

        assert total == brute_total,
               "cost=#{inspect(cost)} gave total #{total}, brute force found #{brute_total}"

        assert valid_assignment?(assignment, n, m),
               "cost=#{inspect(cost)} gave invalid assignment #{inspect(assignment)}"

        assert picked_sum(cost, assignment) == total,
               "cost=#{inspect(cost)}: reported total #{total} != summed picks for #{inspect(assignment)}"
      end
    end
  end

  # An assignment must be an injective mapping row -> column: length n, every
  # value a distinct column in range.
  defp valid_assignment?(assignment, n, m) do
    length(assignment) == n and
      Enum.all?(assignment, &(&1 >= 0 and &1 < m)) and
      length(Enum.uniq(assignment)) == n
  end

  defp picked_sum(cost, assignment) do
    assignment
    |> Enum.with_index()
    |> Enum.reduce(0, fn {col, row}, acc -> acc + (cost |> Enum.at(row) |> Enum.at(col)) end)
  end

  # Exhaustive optimum over every injective row->column mapping.
  defp brute_force(cost) do
    n = length(cost)
    m = length(hd(cost))
    cols = Enum.to_list(0..(m - 1))

    cols
    |> injective_maps(n)
    |> Enum.map(fn assignment -> {picked_sum(cost, assignment), assignment} end)
    |> Enum.min_by(&elem(&1, 0))
  end

  defp injective_maps(_cols, 0), do: [[]]

  defp injective_maps(cols, k) do
    for col <- cols, rest <- injective_maps(cols -- [col], k - 1), do: [col | rest]
  end
end