defmodule MontyExperiments.First.NilMotor do
  use GenServer
  # require Logger

  @default_state %{
    percepts: %{},
    goals: %{},
    state: %{},
  }

  def start_link(id, opts \\ []) do
    GenServer.start_link(__MODULE__, Map.put(@default_state, :id, id), opts)
  end

  def goal_state(motor, id, goal_state) do
    GenServer.cast(motor, {:cortical, self(), id, goal_state})
  end

  def subcortical(motor, id, percept) do
    GenServer.cast(motor, {:subcortical, self(), id, percept})
  end

  def init(data) do
    {:ok, data}
  end

  def handle_call(msg, _from, data) do
    {:reply, {:error, {:unsupported, msg}}, data}
  end

  def handle_cast({:cortical, _pid, id, goal_state}, data) do
    _ = IO.puts("cortial -> id: #{id}, goal_state: #{goal_state}")
    {:noreply, put_in(data, [:goals, id], goal_state)}
  end
  def handle_cast({:subcortical, _pid, id, percept}, data) do
    _ = IO.puts("subcortial -> id: #{id}, percept: #{percept}")
    {:noreply, put_in(data, [:percepts, id], percept)}
  end
  def handle_cast(msg, data) do
    {:noreply, data}
  end

  def handle_info(msg, data) do
    {:noreply, data}
  end
end

defmodule MontyExperiments.First.NilLearning do
  use GenServer
  # require Logger
  alias MontyExperiments.First.NilMotor

  @default_state %{
    subscribers: %{
      output: %{},
      vote: %{},
      bias: %{},
      motor: %{},
    },
    votes: %{},
    bias: %{},
    goals: %{},
    state: %{},
    current_guess: nil,
    current_goal: nil,
  }

  def start_link(id, opts \\ []) do
    GenServer.start_link(__MODULE__, Map.put(@default_state, :id, id), opts)
  end

  def get_hypothesis(lm, timeout \\ :infinity) do
    GenServer.call(lm, :get_hypothesis, timeout)
  end

  def subscribe(lm, type, id, subscriber) do
    GenServer.cast(lm, {:subscribe, type, id, subscriber})
  end

  def unsubscribe(lm, type, id) do
    GenServer.cast(lm, {:unsubscribe, type, id})
  end

  def learn(lm, id, input) do
    GenServer.cast(lm, {:learn, self(), id, input})
  end

  def vote(lm, id, vote) do
    GenServer.cast(lm, {:vote, self(), id, vote})
  end

  def bias(lm, id, bias, goal) do
    GenServer.cast(lm, {:bias, self(), id, bias, goal})
  end

  def init(data) do
    {:ok, data}
  end

  def handle_call(:get_hypothesis, _from, data) do
    {:reply, {:ok, data.current_guess}, data}
  end
  def handle_call(msg, _from, data) do
    {:reply, {:error, {:unsupported, msg}}, data}
  end

  def handle_cast({:subscribe, type, id, subscriber}, data) do
    {:noreply, put_in(data, [:subscribers, type, id], subscriber)}
  end
  def handle_cast({:unsubscribe, type, id}, data) do
    {:noreply, update_in(data, [:subscribers, type], &Map.delete(&1, id))}
  end
  def handle_cast({:bias, _pid, id, cmp, gs}, data) do
    data = data |> put_in([:bias, id], cmp) |> put_in([:goals, id], gs)
    {:noreply, data}
  end
  def handle_cast({:vote, _pid, id, vote}, data) do
    {:noreply, put_in(data, [:votes, id], vote)}
  end
  def handle_cast({:learn, pid, id, cmp}, data) do
    # _ = Logger.debug(fn -> %{msg: :learning, pid: pid, id: id, cmp: cmp, data: data} end)
    # guess
    {:hypothesis, guess, state} = handle_learning(cmp, data.bias, data.state)

    # publish guess to "voting block" LMs
    :ok = Enum.each(data.subscribers.vote, fn {_, lm} -> vote(lm, data.id, guess) end)

    # generate goal
    {:goal, goal, state} = handle_goal_generation(guess, data.goals, state)

    # apply vote
    guess = apply_votes(guess, data.votes)

    # publish goal to Motor Systems
    :ok = Enum.each(data.subscribers.motor, fn {_, ms} -> NilMotor.goal_state(ms, data.id, goal) end)
    # publish bias to lower-level LMs
    :ok = Enum.each(data.subscribers.bias, fn {_, lm} -> bias(lm, data.id, guess, goal) end)
    # publish guess to higher-level LMs
    :ok = Enum.each(data.subscribers.output, fn {_, lm} -> learn(lm, data.id, guess) end)

    {:noreply, %{data | state: state, current_guess: guess, current_goal: goal}}
  end
  def handle_cast(msg, data) do
    {:noreply, data}
  end

  def handle_info(msg, data) do
    {:noreply, data}
  end

  def handle_learning(cmp, biases, state) do
    {:hypothesis, average(cmp, biases), state}
  end

  def handle_goal_generation(guess, goals, state) do
    goals = Enum.map(goals, &elem(&1, 1))
    {:goal, Enum.random([guess | goals]), state}
  end

  defp apply_votes(guess, votes), do: average(guess, votes)

  defp average(guess, other_guesses) do
    (Enum.sum(Enum.map(other_guesses, &elem(&1, 1))) + guess) / (Enum.count(other_guesses) + 1)
  end
end

defmodule MontyExperiments.First.JitterSensor do
  use GenServer
  # require Logger
  alias MontyExperiments.First.NilMotor
  alias MontyExperiments.First.NilLearning

  @default_state %{
    subscribers: %{
      learn: %{},
      motor: %{},
    },
    delay_ms: 100,
    count: 0,
    status: :off
  }

  def start_link(id, opts \\ []) do
    GenServer.start_link(__MODULE__, Map.put(@default_state, :id, id), opts)
  end

  def subscribe(sensor, type, id, subscriber) do
    GenServer.cast(sensor, {:subscribe, type, id, subscriber})
  end

  def unsubscribe(sensor, type, id) do
    GenServer.cast(sensor, {:unsubscribe, type, id})
  end

  def on(sensor), do: GenServer.cast(sensor, :on)
  def off(sensor), do: GenServer.cast(sensor, :off)

  def init(data) do
    _ = schedule_tick(data)
    {:ok, data}
  end

  def handle_call(msg, _from, data) do
    {:reply, {:error, {:unsupported, msg}}, data}
  end

  def handle_cast(:on, %{status: :off} = data) do
    _ = schedule_tick(data)
    {:noreply, %{data | status: :on}}
  end
  def handle_cast(:on, data) do
    {:noreply, data}
  end
  def handle_cast(:off, data) do
    {:noreply, %{data | status: :off}}
  end
  def handle_cast({:subscribe, type, id, subscriber}, data) do
    {:noreply, put_in(data, [:subscribers, type, id], subscriber)}
  end
  def handle_cast({:unsubscribe, type, id}, data) do
    {:noreply, update_in(data, [:subscribers, type], &Map.delete(&1, id))}
  end
  def handle_cast(msg, data) do
    {:noreply, data}
  end

  def handle_info(:tick, %{status: :off} = data) do
    {:noreply, data}
  end
  def handle_info(:tick, %{status: :on} = data) do
    data = tick(data)
    _ = schedule_tick(data)
    {:noreply, data}
  end
  def handle_info(msg, data) do
    {:noreply, data}
  end

  defp tick(%{count: count} = data) do
    :ok = Enum.each(data.subscribers.learn, fn {_, pid} -> NilLearning.learn(pid, data.id, data.count) end)
    :ok = Enum.each(data.subscribers.motor, fn {_, pid} -> NilMotor.subcortical(pid, data.id, data.count) end)
    %{data | count: count + 1}
  end

  defp schedule_tick(data) do
    Process.send_after(self(), :tick, calculate_next_tick(data))
  end

  defp calculate_next_tick(%{delay_ms: delay_ms}) do
    round(delay_ms + 0.5 * :rand.uniform(delay_ms))
  end
end

defmodule MontyExperiments.First.Experiment do
  alias MontyExperiments.First.{NilMotor, NilLearning, JitterSensor}

  # NOTE: use :digraph to describe this instead
  def config_attempt1 do
    %{
      delays: 1..5,
      base_delay: 100,
      sensors: [:s01, :s02, :s03],
      motors: [:m01, :m02],
      learning: [:l01, :l02, :l03, :l11, :l12, :l13, :l21, :l22, :l23],
      sensor_conns: [
        {:s01, :learn, :l01},
        {:s02, :learn, :l02},
        {:s03, :learn, :l03},
        {:s01, :motor, :m01},
        {:s03, :motor, :m02},
      ],
      learning_conns: [
        {:l01, :output, :l11},
        {:l01, :vote, :l02},
        {:l01, :vote, :l03},
        {:l01, :motor, :m01},

        {:l02, :output, :l12},
        {:l02, :vote, :l01},
        {:l02, :vote, :l03},
        {:l02, :motor, :m02},

        {:l03, :output, :l13},
        {:l03, :vote, :l01},
        {:l03, :vote, :l02},
        {:l03, :motor, :m02},

        {:l11, :output, :l21},
        {:l11, :bias, :l01},
        {:l11, :vote, :l12},
        {:l11, :vote, :l13},
        {:l11, :motor, :m01},

        {:l12, :output, :l22},
        {:l12, :bias, :l02},
        {:l12, :vote, :l11},
        {:l12, :vote, :l13},
        {:l12, :motor, :m02},

        {:l13, :output, :l23},
        {:l13, :bias, :l03},
        {:l13, :vote, :l11},
        {:l13, :vote, :l12},
        {:l13, :motor, :m02},

        {:l21, :bias, :l11},
        {:l21, :vote, :l22},
        {:l21, :vote, :l23},
        {:l21, :motor, :m01},

        {:l22, :bias, :l12},
        {:l22, :vote, :l21},
        {:l22, :vote, :l23},
        {:l22, :motor, :m02},

        {:l23, :bias, :l13},
        {:l23, :vote, :l21},
        {:l23, :vote, :l22},
        {:l23, :motor, :m02},
      ],
    }
  end

  def setup(config) do
    sensors = config.sensors |> Enum.map(&{&1, JitterSensor.start_link(&1)}) |> Enum.map(fn {id, {:ok, pid}} -> {id, pid} end) |> Enum.into(%{})
    motors = config.motors |> Enum.map(&{&1, NilMotor.start_link(&1)}) |> Enum.map(fn {id, {:ok, pid}} -> {id, pid} end) |> Enum.into(%{})
    learning = config.learning |> Enum.map(&{&1, NilLearning.start_link(&1)}) |> Enum.map(fn {id, {:ok, pid}} -> {id, pid} end) |> Enum.into(%{})

    :ok = Enum.each(config.sensor_conns, fn
      {sid, :learn, id} -> JitterSensor.subscribe(sensors[sid], :learn, id, learning[id])
      {sid, :motor, id} -> JitterSensor.subscribe(sensors[sid], :motor, id, motors[id])
    end)
    :ok = Enum.each(config.learning_conns, fn
      {lid, :motor, id} -> NilLearning.subscribe(learning[lid], :motor, id, motors[id])
      {lid, type, id} -> NilLearning.subscribe(learning[lid], type, id, learning[id])
    end)

    %{config |
      sensors: sensors,
      motors: motors,
      learning: learning,
    }
  end

  def teardown(_config) do
  end

  def run(config) do
    _ = start_all(config)

    config.delays
    |> Enum.each(fn d ->
      delay = d * config.base_delay
      _ = IO.puts("sleeping for #{delay}ms...")
      :ok = Process.sleep(delay)
      IO.inspect(get_hypotheses(config))
    end)

    _ = stop_all(config)
    config
  end

  def start_all(config) do
    Enum.each(config.sensors, fn {_, pid} -> JitterSensor.on(pid) end)
  end

  def get_hypotheses(config) do
    config.learning
    |> Enum.map(fn {id, pid} ->
      {:ok, hypothesis} = NilLearning.get_hypothesis(pid)
      {id, hypothesis}
    end)
    |> Enum.into(%{})
  end

  def stop_all(config) do
    Enum.each(config.sensors, fn {_, pid} -> JitterSensor.off(pid) end)
  end
end

alias MontyExperiments.First.Experiment

Experiment.config_attempt1()
|> Experiment.setup()
|> Experiment.run()