package aoc.day17

import aoc._
import aoc.direction.Dir
import scala.collection.mutable

import Dir._

extension (d: Dir)
  inline def times(n: Int)(x: Int, y: Int) =
    val k = d.ordinal
    val (nx, ny) = (x + n * dx(k), y + n * dy(k))
    if inside(nx, ny) then
      val dist = d match
        case Up    => sumCols(x)(y) - sumCols(nx)(y)
        case Right => sumRows(x)(ny + 1) - sumRows(x)(y + 1)
        case Down  => sumCols(nx + 1)(y) - sumCols(x + 1)(y)
        case Left  => sumRows(x)(y) - sumRows(x)(ny)
      Some(State(nx, ny, d), dist)
    else None

val board = lines.map(Parser.parse(Parser.row, _).get).map(_.toArray).toArray

val sumRows = board.map(_.scanLeft(0)(_ + _))
val sumCols = (0 until board.length)
  .scanLeft(new Array[Int](board(0).length)) { (last, row) =>
    last.toIterator.zipWithIndex.map { (last, col) => last + board(row)(col) }.toArray
  }
  .toArray

object Parser extends CommonParser:
  val digit = """\d""".r ^^ { _.toInt }
  val row = rep1(digit)
end Parser

inline def inside(x: Int, y: Int) =
  x >= 0 && x < board.length && y >= 0 && y < board(0).length

case class State(x: Int, y: Int, dir: Dir):
  inline def next(ndir: Dir)(using steps: Seq[Int]) =
    steps.toIterator
      .flatMap(ndir.times(_)(x, y))

  inline def left(using Seq[Int]) = next(Dir.fromOrdinal((dir.ordinal + 3) % 4))
  inline def right(using Seq[Int]) = next(Dir.fromOrdinal((dir.ordinal + 1) % 4))

  inline def all(using Seq[Int]) = left ++ right

object Walk:
  // private val visited = mutable.Map.empty[State, Int]
  private object visited:
    private inline def idx(x: Int, y: Int, dir: Dir) = x * board(0).length * 4 + y * 4 + dir.ordinal
    private val arr = Array.fill(board.length * board(0).length * 4)(1_000_000_000)

    inline def apply(s: State) = arr(idx(s.x, s.y, s.dir))
    inline def +=(inline p: (State, Int)) = p match
      case (s, v) => arr(idx(s.x, s.y, s.dir)) = v

  // private val q =
  //   mutable.PriorityQueue.empty[(Int, State)](using scala.math.Ordering.by((v, _) => -v))

  private object q:
    private inline val sz = 128
    private val arrs = mutable.ArrayBuffer.fill(sz)(mutable.ArrayBuffer.empty[State])
    private var ptr = 0
    private var size = 0

    def +=(p: (Int, State)): this.type =
      val (d, s) = p
      size += 1
      arrs(d % sz).addOne(s)
      this

    inline def isEmpty = size == 0

    inline def dequeue() =
      size -= 1
      while arrs(ptr % sz).isEmpty do ptr += 1
      (ptr, arrs(ptr % sz).remove(arrs(ptr % sz).size - 1))

  def go()(using Seq[Int]): Option[Int] =
    Seq(Up, Down, Left, Right).foreach: dir =>
      visited += (State(0, 0, dir) -> 0)
      q += ((0, State(0, 0, dir)))

    @scala.annotation.tailrec
    def loop(): Option[Int] =
      if q.isEmpty then None
      else
        val (d, s) = q.dequeue()
        if d != visited(s) then loop()
        else if s.x == board.length - 1 && s.y == board(0).length - 1 then Some(d)
        else
          for (ns, delta) <- s.all do
            val nd = d + delta
            val old = visited(ns)
            if old > nd then
              visited += (ns -> nd)
              q += ((nd, ns))
          loop()
    var res = loop()
    res

// part 1

def solve(using steps: Seq[Int]) =
  val res =
    Walk
      .go()
      .get
  println(res)

def part1 = solve(using (1 to 3))

def part2 = solve(using (4 to 10))

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