defmodule Aoc2023.Day14 do
  use Aoc2023

  def parse(input) do
    input
    |> lines
    |> Enum.map(&String.to_charlist/1)
  end

  defp replace_many_at(lst, idxs, val) do
    idxs |> Enum.reduce(lst, &List.replace_at(&2, &1, val))
  end

  defp slide(grid) do
    iter = fn recur, grid ->
      case grid do
        [last | [next | rest]] ->
          idxs =
            last
            |> Enum.with_index()
            |> Enum.filter(pipe(elem(0) |> Kernel.==(?O)))
            |> Enum.map(pipe(elem(1)))
            |> Enum.filter(&(Enum.at(next, &1) == ?.))

          [
            last |> replace_many_at(idxs, ?.)
            | recur.(recur, [next |> replace_many_at(idxs, ?O) | rest])
          ]

        [last] ->
          [last]
      end
    end

    loop = fn recur, grid ->
      grid_ = iter.(iter, grid)
      if grid_ == grid, do: grid, else: recur.(recur, grid_)
    end

    loop.(loop, grid |> Enum.reverse())
  end

  defp cycle(grid) do
    for _ <- 1..4, reduce: grid do
      # reverse |> transpose is rotation clockwise, but slide already reverses
      grid -> grid |> slide |> transpose
    end
  end

  defp cycles(grid, n \\ 1_000_000_000, {cache_to_n, cache_from_n} \\ {%{}, %{}}) when n > 0 do
    rep_n = Map.get(cache_to_n, grid)
    if rep_n do
      n_ = rep_n - Integer.mod(n, rep_n - n)
      cache_from_n |> Map.get(n_)
    else
      grid_ = cycle(grid)
      cycles(grid_, n - 1, {Map.put(cache_to_n, grid, n), Map.put(cache_from_n, n, grid)})
    end
  end

  def part1(input) do
    input
    |> slide
    |> Enum.with_index(1)
    |> Enum.map(fn {row, i} -> Enum.count(row, &(&1 == ?O)) * i end)
    |> Enum.sum()
  end

  def part2(input) do
    input
    |> cycles()
    |> Enum.reverse
    |> Enum.with_index(1)
    |> Enum.map(fn {row, i} -> Enum.count(row, &(&1 == ?O)) * i end)
    |> Enum.sum()
  end
end