defmodule Aoc2024.Day08 do
  def parse(input) do
    input = input |> String.trim() |> String.split("\n") |> Enum.map(&String.codepoints/1)
    w = Enum.count(Enum.at(input, 0))
    h = Enum.count(input)

    input =
      walk(input, w, h)
      |> Enum.reduce(%{}, fn {x, i, j}, m ->
        if Map.has_key?(m, x) do
          %{m | x => [{i, j} | m[x]]}
        else
          Map.put(m, x, [{i, j}])
        end
      end)

    {input, w, h}
  end

  defp walk(t, w, h), do: walk(t, 0, 0, w, h, [])
  defp walk(_t, i, _j, _w, h, towers) when i == h, do: towers
  defp walk(t, i, j, w, h, towers) when j == w, do: walk(t, i + 1, 0, w, h, towers)

  defp walk(t, i, j, w, h, towers) do
    case t |> Enum.at(i) |> Enum.at(j) do
      "." -> walk(t, i, j + 1, w, h, towers)
      x -> walk(t, i, j + 1, w, h, [{x, i, j} | towers])
    end
  end

  def part1({input, w, h}) do
    for {_x, all} <- input do
      Task.async(fn -> comb(&antinodes/4, all) end)
    end
    |> Enum.map(&Task.await/1)
    |> List.flatten()
    |> Enum.into(MapSet.new())
    |> Enum.filter(fn {i, j} -> i >= 0 and j >= 0 and i < h and j < w end)
    |> Enum.count()
  end

  defp comb(antinodes, [current | others]), do: comb(antinodes, others, current, others, [])
  defp comb(_antinodes, [_], _, [], res), do: res

  defp comb(antinodes, [current | others], _, [], res),
    do: comb(antinodes, others, current, others, res)

  defp comb(antinodes, others, {i1, j1} = current, [{i2, j2} | tail], res) do
    comb(antinodes, others, current, tail, [antinodes.(i1, j1, i2, j2) | res])
  end

  defp antinodes(i1, j1, i2, j2) do
    di = i2 - i1
    dj = j2 - j1
    [{i1 - di, j1 - dj}, {i2 + di, j2 + dj}]
  end

  def part2({input, w, h}) do
    for {_x, all} <- input do
      Task.async(fn -> comb(fn i1, j1, i2, j2 -> antinodes2(w, h, i1, j1, i2, j2) end, all) end)
    end
    |> Enum.map(&Task.await/1)
    |> List.flatten()
    |> Enum.into(MapSet.new())
    |> Enum.count()
  end

  defp antinodes2(w, h, i1, j1, i2, j2) do
    di = i2 - i1
    dj = j2 - j1

    Stream.iterate({i1, j1}, fn {i, j} -> {i + di, j + dj} end)
    |> Stream.take_while(fn {i, j} -> i >= 0 and j >= 0 and i < h and j < w end)
    |> Stream.concat(
      Stream.iterate({i2, j2}, fn {i, j} -> {i - di, j - dj} end)
      |> Stream.take_while(fn {i, j} -> i >= 0 and j >= 0 and i < h and j < w end)
    )
    |> Enum.into([])
  end
end