package aoc.day14

import aoc._
import scala.collection.mutable

type Board = Seq[Seq[Char]]

val board = lines.map(_.toSeq).toSeq

enum Direction { case Up, Down, Left, Right }
import Direction._

extension (board: Board)
  inline def doTiltLeft(inline before: Board => Board, inline after: Board => Board): Board =
    after(before(board).map { col =>
      col
        .splitBy(_ == '#') // split by tilted rocks
        .map { group =>
          val (rocks, empty) = group.partition(_ == 'O')
          (rocks ++ empty).mkString
        }
        .mkString("#")
        .toSeq
    })

  def doTilt(direction: Direction): Board = direction match
    case Up    => doTiltLeft(_.transpose, _.transpose)
    case Down  => doTiltLeft(_.transpose.map(_.reverse), _.map(_.reverse).transpose)
    case Left  => doTiltLeft(identity, identity)
    case Right => doTiltLeft(_.map(_.reverse), _.map(_.reverse))

  def cycle = Seq(Up, Left, Down, Right).foldLeft(board)(_.doTilt(_))

  def load =
    board.reverse // rows in reverse
      .zipWithIndex.map { case (row, idx) => row.count(_ == 'O') * (idx + 1L) }.sum

// part 1

def part1 =
  println(board.doTilt(Up).load)

// part 2

object Cycling:
  private val memo = mutable.Map.empty[(Board, Long), Board]

  def apply(b: Board, times: Long): Board =
    if times == 0 then b
    else
      memo.getOrElseUpdate(
        (b, times), {
          val twice = apply(apply(b, times / 2), times / 2)
          if times % 2 == 0 then twice else twice.cycle
        }
      )

def part2 =
  println(Cycling(board, 1_000_000_000).load)

@main def Day14(part: Int) = part match
  case 1 => part1
  case 2 => part2