package aoc.day4

import aoc._
import scala.collection.mutable.Map

case class Card(id: Int, winning: List[Int], got: List[Int]):
  lazy val won =
    val winSet = winning.toSet
    got.filter(winSet contains _)
  def pointValue =
    if won.isEmpty then 0 else 1L << (won.length - 1)

object Parser extends CommonParser:
  val nums = rep1(num)
  val card = "Card " ~ num ~ ":" ~ nums ~ "|" ~ nums ^^ { case (_ ~ id ~ _ ~ winning ~ _ ~ got) =>
    Card(id, winning, got)
  }
end Parser

val cards = lines.map(l => Parser.parse(Parser.card, l).get).toSeq

// Part 1

def part1 =
  val res = cards.map(_.pointValue).sum
  println(res)

// Part 2

class Tails():
  private val map = Map[Int, Long]()
  var sum = 0L

  def add(tail: Int, copies: Long): Unit =
    sum += copies
    map.updateWith(tail)(v => Some(v.fold(copies)(_ + copies)))

  def filter(id: Int): Unit =
    map.remove(id).foreach { sum -= _ }
end Tails

@scala.annotation.tailrec
def traverse(ids: Iterator[Card], tails: Tails, accum: Long = 0): Long =
  if !ids.hasNext then return accum
  val card = ids.next()
  val copies = 1 + tails.sum
  tails.add(card.id + card.won.length, copies)
  tails.filter(card.id)
  traverse(ids, tails, accum + copies)

def part2 =
  val res = traverse(cards.toIterator, new Tails())
  println(res)

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