package aoc.day22

import aoc._
import scala.collection.mutable

case class Point3(x: Int, y: Int, z: Int):
  def min(other: Point3) = Point3(x.min(other.x), y.min(other.y), z.min(other.z))
  def max(other: Point3) = Point3(x.max(other.x), y.max(other.y), z.max(other.z))

  def to2D = (x, y)

case class Brick private (index: Int, low: Point3, high: Point3):
  inline def height = high.z - low.z + 1
  def to2D = (low.to2D, high.to2D)

object Brick:
  private var index = 0
  def fromPoints(a: Point3, b: Point3) =
    index += 1
    Brick(index, a.min(b), a.max(b))

object Parser extends CommonParser:
  val point3 = num ~ "," ~ num ~ "," ~ num ^^ { case x ~ _ ~ y ~ _ ~ z => Point3(x, y, z) }
  val brick = point3 ~ "~" ~ point3 ^^ { case a ~ _ ~ b => Brick.fromPoints(a, b) }

val bricks =
  lines
    .map(Parser.parse(Parser.brick, _).get)
    .toSeq

// part 1

class Space:
  case class Cell(height: Int, occupied: Int /* index of the occupied brick */ )
  val space = mutable.Map.empty[(Int, Int), Cell]

  // Override a brick falling and returns the height and unique brick that holds it, if it exists.
  def overrideBrick(b: Brick): (Int, Option[Int]) =
    // the space covered is
    val covered = (b.low.x to b.high.x).flatMap(x => (b.low.y to b.high.y).map((x, _)))
    // get the maximum height occupied
    val pastCells = covered.flatMap(space.get(_))
    val maxHeight = pastCells.map(_.height).maxOption.getOrElse(0)
    assert(maxHeight <= b.low.z)
    // update space
    covered.foreach(space.update(_, Cell(maxHeight + b.height, b.index)))
    // get the ones held
    val res = pastCells.filter(_.height == maxHeight).map(_.occupied).toSet.toSeq match
      case Seq(idx) => Some(idx)
      case _        => None
    (maxHeight, res)

// sort bricks by low z
val sortedBricks = bricks.sortBy(_.low.z)

def part1() =
  val space = Space()
  val holdingBricks = sortedBricks.foldLeft(Set.empty[Int])(_ ++ space.overrideBrick(_)._2)
  println(bricks.size - holdingBricks.size)

def part2() =
  extension (bs: Iterator[Brick])
    def heights =
      val space = Space()
      bs.map(b => (b -> space.overrideBrick(b)._1)).toMap
  extension (bm: Map[Brick, Int])
    inline def countDifferences(original: Map[Brick, Int]) =
      bm.count((b, h) => h != original(b))

  val original = sortedBricks.iterator.heights
  val res = bricks
    .map: exclude =>
      sortedBricks.iterator
        .filter(_ != exclude)
        .heights
        .countDifferences(original)
    .sum
  println(res)

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