defmodule MbBinaryTree.Traverse do
  @moduledoc """
  Utility functions for traversing binary trees.
  """

  @doc """
  Perform a preorder traversal of a binary tree.

  ## Examples

      iex> MbBinaryTree.Traverse.preorder([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list()
      [1, 2, 3]
      iex> MbBinaryTree.Traverse.preorder([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list()
      [12, 4, 2, 5, 26, 14, 28, 27, 30]
  """
  def preorder(tree) do
    Stream.unfold([tree], &preorder_acc/1)
  end

  defp preorder_acc([tree | stack]) do
    case tree do
      [value | [left, right]] -> {value, [left, right | stack]}
      nil -> preorder_acc(stack)
    end
  end

  defp preorder_acc([]), do: nil

  @doc """
  Perform a postoder traversal of a binary tree.

  ## Examples

      iex> MbBinaryTree.Traverse.postorder([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list()
      [2, 3, 1]
      iex> MbBinaryTree.Traverse.postorder([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list()
      [2, 5, 4, 14, 27, 30, 28, 26, 12]
  """
  def postorder(tree) do
    Stream.unfold([tree], &postorder_acc/1)
  end

  defp postorder_acc([tree | stack]) do
    case tree do
      [value | [nil, nil]] -> {value, stack}
      [value | [left, nil]] -> postorder_acc([left, [value, nil, nil] | stack])
      [value | [nil, right]] -> postorder_acc([right, [value, nil, nil] | stack])
      [value | [left, right]] -> postorder_acc([left, right, [value, nil, nil] | stack])
      nil -> postorder_acc(stack)
    end
  end

  defp postorder_acc([]), do: nil

  @doc """
  Perform a breadth-first traversal of a binary tree.

  ## Examples

      iex> MbBinaryTree.Traverse.breadth_first([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list()
      [1, 2, 3]
      iex> MbBinaryTree.Traverse.breadth_first([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list()
      [12, 4, 26, 2, 5, 14, 28, 27, 30]
  """
  def breadth_first(tree) do
    queue = :queue.from_list([tree])
    Stream.unfold(queue, &breadth_first_acc/1)
  end

  defp breadth_first_acc({[], []}), do: nil

  defp breadth_first_acc(queue) do
    {{:value, tree}, queue} = :queue.out(queue)

    case tree do
      [value | [nil, nil]] -> {value, queue}
      [value | [left, nil]] -> {value, queue |> q_in(left)}
      [value | [nil, right]] -> {value, queue |> q_in(right)}
      [value | [left, right]] -> {value, queue |> q_in(left) |> q_in(right)}
      nil -> breadth_first_acc(queue)
    end
  end

  # A convenience function to allow piping above by swapping the argument order.
  defp q_in(queue, item), do: :queue.in(item, queue)

  @doc """
  Perform an inorder traversal of a binary tree.

  ## Examples

      iex> MbBinaryTree.Traverse.inorder([1, [2, nil, nil], [3, nil, nil]]) |> Enum.to_list()
      [2, 1, 3]
      iex> MbBinaryTree.Traverse.inorder([12, [4, MbBinaryTree.new(2), MbBinaryTree.new(5)], [26, MbBinaryTree.new(14), [28, MbBinaryTree.new(27), MbBinaryTree.new(30)]]]) |> Enum.to_list()
      [2, 4, 5, 12, 14, 26, 27, 28, 30]
  """
  def inorder(tree) do
    Stream.unfold([tree], &inorder_acc/1)
  end

  defp inorder_acc([tree | stack]) do
    case tree do
      [value | [nil, nil]] -> {value, stack}
      [value | [left, nil]] -> inorder_acc([left, [value | [nil, nil]] | stack])
      [value | [nil, right]] -> {value, [right | stack]}
      [value | [left, right]] -> inorder_acc([left, [value | [nil, nil]], right | stack])
      nil -> inorder_acc(stack)
    end
  end

  defp inorder_acc([]), do: nil
end