defmodule DistNeatEx.NodeKnower do
  use GenServer
  alias DistNeatEx.NodeKnower
  defstruct species: %{}, nodes: %{}, timestamps: %{}

  ## state structure ##
  # %NodeKnower{
  #   species: %{
  #     <ref>: %SpeciesKnowledge{},
  #     ...
  #   },
  #   nodes: %{
  #     <node_name>: %{
  #       <run_id_term>: MapSet<[<ref1>, <ref2>, ...]>,
  #       ...
  #     },
  #     ...
  #   },
  #   timestamps: %{ TODO: Use this to check species health
  #     <ref>: <timestamp>,
  #     ...
  #   }
  # }

  defmodule SpeciesKnowledge do
    @moduledoc "[:node_name, :run_id_term, :representative, :avg_fitness, :generation]"
    defstruct [:node_name, :run_id_term, :representative, :avg_fitness, :best_ann, :generation]

    @doc "Requires data %{:node_name, :run_id_term, :representative, :avg_fitness, :generation}"
    def new(data), do: struct!(SpeciesKnowledge, data)
  end

  def start_link(args) do
    GenServer.start_link(__MODULE__, args, name: __MODULE__)
  end

  def init(_args) do
    :net_kernel.monitor_nodes(true) # monitor nodeup and nodedown events
    # Then ensure all current nodes are in NodeKnower
    state = %NodeKnower{ nodes:
      Enum.reduce(all_nodes(), %{}, fn node, nodes ->
        Map.put(nodes, node, %{}) # set to empty map if non-existant
      end)
    }
    {:ok, state}
  end

  def all_nodes(), do: [node() | Node.list()] |> Enum.sort()

  @doc "Calls to the NodeKnower. This ensures the NodeKnower is running before calling."
  def call(msg) do
    GenServer.call(__MODULE__, msg)
  end

  @doc "Casts to the NodeKnower. This ensures the NodeKnower is running before casting."
  def cast(msg) do
    GenServer.abcast(__MODULE__, msg)
    :ok
  end

  @doc "Handles when a node comes up."
  def handle_info({:nodeup, node}, state) do
    state = put_in(state, [Access.key!(:nodes), node], %{})
    {:noreply, state}
  end
  @doc "Handles when a node goes down, and removes its entries from the node_knower"
  def handle_info({:nodedown, node}, state=%NodeKnower{species: species, nodes: nodes}) do
    runs = Map.get(nodes, node, [])
    nodes = Map.delete(nodes, node)
    species = Enum.reduce(runs, species, fn {_run_id_term, ref_set}, species ->
      Enum.reduce(ref_set, species, fn ref, species ->
        Map.delete(species, ref)
      end)
    end)
    {:noreply, %{state | nodes: nodes, species: species}}
  end

  @doc "Update the NodeKnower with the latest species_knowledge, given the species ref."
  def update_species(ref, species_knowledge), do: cast({:update_species, ref, species_knowledge})
  def update_species_helper(state=%NodeKnower{species: species, nodes: nodes, timestamps: timestamps}, ref, sk=%SpeciesKnowledge{}) do
    old_sk = Map.get(species, ref)
    species = Map.put(species, ref, sk)

    nodes = cond do
      is_nil(old_sk) ->
        # Creating a new species
        nodes
        |> put_ref_in_nodes(ref, sk)
      old_sk.node_name != sk.node_name ->
        # Species moved nodes
        nodes
        |> del_ref_from_nodes(ref, old_sk)
        |> put_ref_in_nodes(ref, sk)
      :else ->
        nodes
    end

    timestamps = Map.put(timestamps, ref, :erlang.monotonic_time())

    %{state | species: species, nodes: nodes, timestamps: timestamps}
  end
  def update_species_helper(state, ref, sk), do: update_species_helper(state, ref, SpeciesKnowledge.new(sk))
  defp put_ref_in_nodes(nodes, ref, sk) do
    Map.update!(nodes, Map.fetch!(sk, :node_name), fn runs ->
      Map.update(runs, sk.run_id_term, MapSet.new([ref]), fn mapset ->
        MapSet.put(mapset, ref)
      end)
    end)
  end
  defp del_ref_from_nodes(nodes, ref, sk=%SpeciesKnowledge{}) do
    Map.update!(nodes, Map.fetch!(sk, :node_name), fn runs ->
      Map.update(runs, sk.run_id_term, MapSet.new(), fn mapset ->
        MapSet.delete(mapset, ref)
      end)
    end)
  end
  defp del_ref_from_nodes(nodes, ref, sk), do: del_ref_from_nodes(nodes, ref, SpeciesKnowledge.new(sk))

  def node_num_species(state, node, run_id) do
    state.nodes
    |> Map.fetch!(node)
    |> Map.get(run_id, MapSet.new)
    |> MapSet.size()
  end
  def node_sum_fitness(state=%{species: species}, node, run_id) do
    state.nodes
    |> Map.fetch!(node)
    |> Map.get(run_id, [])
    |> Enum.reduce(0, fn ref, sum ->
      sum + Map.fetch!(species, ref).avg_fitness
    end)
  end
  def node_avg_fitness(state, node, run_id) do
    num = node_num_species(state, node, run_id)
    sum = node_sum_fitness(state, node, run_id)
    if (num == 0), do: :infinity, else: sum / num
  end

  def node_cross_run_sum_fitness(state, node) do
    species = state.species
    Enum.reduce(state.nodes[node] || [], 0, fn {_run_id_term, ref_set}, sum ->
      Enum.reduce(ref_set, sum, fn ref, sum ->
        sum + Map.fetch!(species, ref).avg_fitness
      end)
    end)
  end

  def cross_node_sum_fitness(state, run_id) do
    species = state.species
    Enum.reduce(state.nodes, 0, fn {_node, run_map}, sum ->
      Enum.reduce(run_map[run_id] || [], sum, fn ref, sum ->
        sum + Map.fetch!(species, ref).avg_fitness
      end)
    end)
  end

  def cross_node_num_species(state, run_id) do
    Enum.reduce(state.nodes, 0, fn {_node, run_map}, sum ->
      sum + MapSet.size(run_map[run_id] || %MapSet{})
    end)
  end

  def cross_node_best_ann(state, run_id) do
    cross_node_max_by_key(state, run_id, :best_ann, fn {_ann, fitness} -> fitness end)
  end

  def cross_node_max_by_key(state, run_id, key, comparefn \\ &(&1)) do
    species = state.species
    flat = Enum.map(state.nodes, fn {_node, run_map} ->
      Enum.map(run_map[run_id] || [], fn ref ->
        {ref, Map.fetch!(species, ref) |> Map.fetch!(key)}
      end)
    end)
    |> List.flatten

    if flat == [] do
      {nil, nil}
    else
      Enum.max_by(flat, fn {_, val} -> comparefn.(val) end)
    end
  end

  def get_species_of_run(state, run_id) do
    state.nodes
    |> Enum.flat_map(fn {_, map} ->
      map
      |> Map.get(run_id, MapSet.new())
      |> MapSet.to_list
    end)
  end

  def get_run_ids(state) do
    state.nodes
    |> Enum.reduce(MapSet.new(), fn {_, map}, set ->
      MapSet.union(set, MapSet.new(Map.keys(map)))
    end)
    |> Enum.to_list
  end

  @doc "Calls node knower to get the best ann. Returns a {ann, fitness} pair."
  def best_ann(run_id), do: call({:best_ann, run_id})

  def can_i_die?(run_id), do: call({:can_i_die, run_id})

  def send_me_state(), do: GenServer.cast(__MODULE__, {:send_me_state, self()})

  def species_up(ref, sk_args, neat_opts), do: cast({:species_up, ref, sk_args, neat_opts})
  @doc "If neat_opts is nil, it indicates that the run should not be restarted if this is the last species to go down."
  def species_down(ref, sk, neat_opts), do: cast({:species_down, ref, sk, neat_opts})

  def handle_species_up(state, ref, sk_args, _neat_opts) do
    NodeKnower.update_species_helper(state, ref, sk_args)
  end
  def handle_species_down(state, ref, sk_args, neat_opts) do
    state = state
    |> Map.update!(:species, &Map.delete(&1, ref))
    |> Map.update!(:nodes, &del_ref_from_nodes(&1, ref, sk_args))

    run_id = sk_args[:run_id_term]
    if neat_opts && cross_node_num_species(state, run_id) == 0 do
      IO.puts("!!!!!!!!! Extinction occured, restarting run #{inspect run_id} !!!!!!!!!")
      spawn_link(fn -> DistNeatEx.Helper.init_run(run_id, neat_opts) end)
      state
    else
      state
    end
  end

  @doc "Returns the node with the lowest sum of average fitness"
  def most_available_node(state) do
    all_nodes() #TODO compute this only for our run, so fitness can be independent.
    |> Enum.min_by(&node_cross_run_sum_fitness(state, &1))
  end
  def most_available_node(), do: call(:most_available_node)

  def handle_cast({:update_species, ref, sk}, state), do: {:noreply, update_species_helper(state, ref, sk)}
  def handle_cast({:send_me_state, pid}, state) do
    send(pid, {:node_knower_state, state})
    {:noreply, state}
  end

  def handle_cast({:species_up, ref, sk_args, neat_opts}, state) do
    state = handle_species_up(state, ref, sk_args, neat_opts)
    {:noreply, state}
  end
  def handle_cast({:species_down, ref, sk, neat_opts}, state) do
    state = handle_species_down(state, ref, sk, neat_opts)
    {:noreply, state}
  end

  def handle_cast({:delete_run, run_id_term}, state) do
    refs = get_species_of_run(state, run_id_term)
    state = Enum.reduce(refs, state, fn ref, state ->
      state = update_in(state.species, &Map.delete(&1, ref))
      update_in(state.timestamps, &Map.delete(&1, ref))
    end)
    state = update_in(state.nodes, fn nodes ->
      Enum.reduce(nodes, %{}, fn {node, map}, new_nodes ->
        Map.put(new_nodes, node, Map.delete(map, run_id_term))
      end)
    end)

    {:noreply, state}
  end

  def handle_call(:get_state, _from, state) do
    {:reply, state, state}
  end

  def handle_call(:most_available_node, _from, state) do
    {:reply, most_available_node(state), state}
  end

  def handle_call({:best_ann, run_id}, _from, state) do
    {_ref, best_ann} = cross_node_best_ann(state, run_id)
    {:reply, best_ann, state}
  end

  def handle_call({:can_i_die, run_id}, _from, state) do
    {:reply, cross_node_num_species(state, run_id) > 1, state}
  end
end