Operation Order

defmodule Aoc2020.Day18 do
  @moduledoc "Operation Order"

  def run() do
    part_2()
  end

  def part_1() do
    parse_input()
    |> Enum.map(&evaluate/1)
    |> Enum.sum()
  end

  def part_2() do
    parse_input()
    |> Enum.map(&evaluate2/1)
  end

  def parse_input() do
    File.read!("priv/inputs/2020/day18.txt")
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(&parse_expression/1)
  end

  def parse_expression(str) do
    str
    |> String.codepoints()
    |> (fn tokens -> parse_tokens(tokens) end).()
  end

  def parse_tokens([], tree), do: Enum.reverse(tree)

  def parse_tokens(["(" | rest], tree) do
    {block, rest} = parse_tokens(rest, [])
    parse_tokens(rest, [block | tree])
  end

  def parse_tokens([")" | rest], tree) do
    {Enum.reverse(tree), rest}
  end

  def parse_tokens([head | rest], tree \\ []) do
    case head do
      " " -> parse_tokens(rest, tree)
      "+" -> parse_tokens(rest, ["+" | tree])
      "*" -> parse_tokens(rest, ["*" | tree])
      d -> parse_tokens(rest, [String.to_integer(d) | tree])
    end
  end

  def evaluate([x | []]), do: x

  def evaluate([x, op, y | rest]) do
    val = apply(operation(op), [flatten(x), flatten(y)])
    evaluate([val | rest])
  end

  def operation("+"), do: &+/2
  def operation("*"), do: &*/2

  def evaluate2([x, "+", y | rest]),
    do: [flatten(x, &evaluate2/1) + flatten(y, &evaluate2/1) | rest]

  def evaluate2([x, "*", y | rest]),
    do: [flatten(x, &evaluate2/1) * flatten(y, &evaluate2/1) | rest]

  def flatten(x) when is_integer(x), do: x
  def flatten(x, eval_fn \\ &evaluate/1) when is_list(x), do: apply(eval_fn, [x])
end