defmodule Markovian do
  @moduledoc """
  Markovian is a Q-learner for the BEAM.
  """

  defstruct alpha: 0.2, gamma: 0.9, q_table: %{}, current_state: nil, current_action: nil,
    random_action_rate: 0.4, random_action_decay: 0.99

  @type t :: %__MODULE__{}
  @type state :: term
  @type action :: term

  @doc """
  Creates a new MDP. To build a new MDP, you need to provide a set of states
  and actions as well an initial state.
  """
  @spec new(states::[state], actions::[action], initial_state::state) :: t
  def new(states, actions, initial_state) do
    state_space = for state <- states, into: %{} do
      action_space = for action <- actions, into: %{} do
        {action, 2 * :rand.uniform() - 1}
      end

      {state, action_space}
    end

    %__MODULE__{
      q_table: state_space,
      current_state: initial_state,
      current_action: actions |> Enum.random()
    }
  end

  @doc """
  Updates an MDP. It needs to get the reward and next state from the previous
  state and action.
  """
  @spec update(t, float, state) :: t
  def update(mdp, reward, next_state) do
    s_update = (1 - mdp.alpha) * get_in(mdp.q_table, [mdp.current_state, mdp.current_action])
    next_q = mdp.q_table[next_state] |> Map.values() |> Enum.max()
    s_prime_update = mdp.alpha * (reward + mdp.gamma * next_q)
    updated_q_table = put_in(mdp.q_table, [mdp.current_state, mdp.current_action], s_update + s_prime_update)
    %{mdp | q_table: updated_q_table, current_state: next_state}
    |> generate_action()
    |> decay()
  end

  defp generate_action(mdp) do
    random? = mdp.random_action_rate > :rand.uniform()
    action = do_generate_action(mdp.q_table[mdp.current_state], random?)
    %{mdp | current_action: action}
  end

  defp do_generate_action(current_state_actions, random?)
  defp do_generate_action(current_state_actions, false) do
    current_state_actions
    |> Map.to_list()
    |> Enum.max_by(&elem(&1, 1))
    |> elem(0)
  end
  defp do_generate_action(current_state_actions, true) do
    current_state_actions
    |> Map.keys()
    |> Enum.random()
  end

  defp decay(mdp) do
    Map.update!(mdp, :random_action_rate, &(&1 * mdp.random_action_decay))
  end
end