kmeans

defmodule Kmeans.KmeansCommand do
  defmodule ClosestCentroidCommand do
    def execute(element, centroids, distance_fun) do
      {centroid, _diff} = closest_centroid(element, centroids, distance_fun)
      {element, centroid}
    end

    defp closest_centroid(element, centroids, distance_fun) do
      centroids
      |> Enum.reduce({}, smallest_distance(element, distance_fun))
    end

    defp smallest_distance(element, distance_fun) do
      fn (centroid, memo) ->
        smaller {centroid, distance_fun.(element, centroid)}, memo
      end
    end

    defp smaller(original, {}), do: original
    defp smaller({_, diff} = current, {_, smallest_diff} = memo) do
      smallest current, memo, diff < smallest_diff
    end

    defp smallest(current, _memo, true), do: current
    defp smallest(_current, memo, false), do: memo
  end

  def execute(centroids, _data, 0, _fun1, _fun2), do: centroids
  def execute(centroids, data, loop, distance_fun, mean_fun) do
    data
    |> run_async(centroids, distance_fun)
    |> Enum.to_list
    |> Enum.reduce(%{}, &build_map/2)
    |> include_all_centroids(centroids)
    |> compute_new_centroids(mean_fun)
    |> execute(data, loop - 1, distance_fun, mean_fun)
  end

  defp run_async(list, centroids, distance_fun) do
    list
    |> async_stream(centroids, distance_fun)
    |> Stream.filter(&success?/1)
    |> Stream.map(&transform/1)
  end

  defp async_stream(list, centroids, distance_fun) do
    {:ok, supervisor} = Task.Supervisor.start_link

    supervisor
    |> Task.Supervisor.async_stream_nolink(list,
      ClosestCentroidCommand, :execute, [centroids, distance_fun])
  end

  defp compute_new_centroids(clusters, mean_fun) do
    clusters
    |> Enum.with_index
    |> Enum.reduce([], &(add_centroid(&1, &2, mean_fun)))
  end

  defp add_centroid({{centroid, values}, index}, clusters, fun) do
    centroid = build_centroid(centroid, values, index, fun)
    [centroid | clusters]
  end

  defp build_centroid(centroid, [], index, _fun) do
    centroid_with_id centroid, index
  end
  defp build_centroid(_centroid, values, index, fun) do
    fun.(values)
    |> centroid_with_id(index)
  end

  defp centroid_with_id(centroid, id), do: Map.put centroid, :id, id + 1

  defp include_all_centroids(clusters, centroids) do
    Enum.reduce centroids, clusters, &add_centroid_to_clusters/2
  end

  defp add_centroid_to_clusters(centroid, clusters) do
    add_centroid_to_clusters centroid, clusters, clusters[centroid]
  end
  defp add_centroid_to_clusters(centroid, clusters, nil) do
    Map.put clusters, centroid, []
  end
  defp add_centroid_to_clusters(_centroid, clusters, _add?), do: clusters

  defp success?({:ok, _}), do: true
  defp success?(_), do: false

  defp transform({:ok, value}), do: value

  defp build_map({element, centroid}, memo) do
    add_element memo, centroid, element, memo[centroid]
  end

  defp add_element(memo, centroid, element, nil) do
    memo
    |> Map.put(centroid, [element])
  end
  defp add_element(memo, centroid, element, _) do
    memo
    |> Map.put(centroid, [element | memo[centroid]])
  end
end