defmodule Day12 do
  def part1(data) do
    grid = data
    |> String.split("\n", trim: true)
    |> Enum.map(&String.codepoints/1)

    grid = for {l, i} <- Enum.with_index(grid),
               {c, j} <- Enum.with_index(l) do
             {{i, j}, c}
    end
    |> Map.new()

    for pos <- Map.keys(grid), reduce: {MapSet.new(), 0} do
      {vis, res} -> solve1(vis, grid, pos)
      |> then(fn {vis, a, p} -> {vis, res + a*p} end)
    end
    |> elem(1)
  end

  defp solve1(vis, grid, pos) do
    if MapSet.member?(vis, pos) do
      {vis, 0, 0}
    else
      for npos <- dirs(pos), reduce: {MapSet.put(vis, pos), 1, 0} do
        {vis, a, p} ->
          case Map.fetch(grid, npos) do
            :error ->
              {vis, a, p+1}
            {_, v} ->
              if v == Map.fetch!(grid, pos) do
                {vis, na, np} = solve1(vis, grid, npos)
                {vis, a+na, p+np}
              else
                {vis, a, p+1}
              end
          end
      end
    end
  end

  defp dirs({i, j}) do
    [{i+1, j}, {i-1, j}, {i, j+1}, {i, j-1}]
  end

  def part2(data) do
    grid = data
    |> String.split("\n", trim: true)
    |> Enum.map(&String.codepoints/1)

    grid = for {l, i} <- Enum.with_index(grid),
               {c, j} <- Enum.with_index(l) do
             {{i, j}, c}
    end
    |> Map.new()

    for pos <- Map.keys(grid), reduce: {MapSet.new(), 0} do
      {vis, res} -> solve2(vis, grid, pos)
      |> then(fn {vis, a, walls} -> {vis, res + a*perim(walls)} end)
    end
    |> elem(1)
  end

  defp solve2(vis, grid, pos) do
    if MapSet.member?(vis, pos) do
      {vis, 0, []}
    else
      for {npos, dir} <- dirs2(pos), reduce: {MapSet.put(vis, pos), 1, []} do
        {vis, a, walls} ->
          case Map.fetch(grid, npos) do
            :error ->
              {vis, a, [{dir, npos} | walls]}
            {_, v} ->
              if v == Map.fetch!(grid, pos) do
                {vis, na, nwalls} = solve2(vis, grid, npos)
                {vis, a+na, walls ++ nwalls}
              else
                {vis, a, [{dir, npos} | walls]}
              end
          end
      end
    end
  end

  defp dirs2({i, j}) do
    [{{i+1, j}, :down},
     {{i-1, j}, :up},
     {{i, j+1}, :right},
     {{i, j-1}, :left}]
  end

  defp dist({a, b}, {c, d}), do: abs(a-c) + abs(b-d)
  defp dist(_, nil), do: nil
  defp dist(nil, _), do: nil

  defp perim([]), do: 0

  defp perim(walls) do
    gwalls = walls
    |> Enum.group_by(&elem(&1,0), &elem(&1,1))

    for {dir, walls} <- gwalls do
      swalls =
        cond do
          dir == :up or dir == :down ->
            Enum.sort(walls)
          true ->
            Enum.sort_by(walls, fn {i, j} -> {j, i} end)
        end

      for w <- swalls, reduce: {nil, 0} do
        {prev, acc} ->
          if dist(w, prev) == 1 do
            {w, acc}
          else
            {w, acc+1}
          end
      end
      |> elem(1)
    end
    |> Enum.sum()
  end
end

data = IO.read(:stdio, :eof)

{time1 , ans1} = :timer.tc(fn -> Day12.part1(data) end)
IO.puts("Time  : #{time1 / 1000000}")
IO.puts("Answer: #{ans1}")

{time2 , ans2} = :timer.tc(fn -> Day12.part2(data) end)
IO.puts("Time  : #{time2 / 1000000}")
IO.puts("Answer: #{ans2}")