Paste: aoc24 p1

TUPLE: armygroup army id units hp attack initiative attack-type immunes weaknesses ;
C: <armygroup> armygroup

: parse-parens ( str -- immunes/f weakness/f )
"()" split dup length 3 = [
  second ";" split [ [ CHAR: space = ] trim-head ] map
  [ " " split 1 swap remove-nth [ [ CHAR: , = ] trim-tail ] map ] map
  dup length 1 =
  [ dup first first "immune" = [ f suffix ] [ f prefix ] if ]
  [ dup first first "weak" = [ reverse ] when ] if
  [ [ rest ] [ f ] if* ] map first2
] [ drop f f ] if ;
: parse-nums ( str -- units hp attack initiative attack-type )
" " split
[ { 0 4 } swap nths ]
[ 6 tail* { 0 5 } swap nths append [ string>number ] map first4 ] 
[ 5 tail* first ] tri ;

: effective-power ( group -- power ) [ units>> ] [ attack>> ] bi * ;
: damage ( group1 group2 -- damage )
  [ [ effective-power ] [ attack-type>> ] bi ]
  [ [ immunes>> ] [ weaknesses>> ] bi ] bi*
  {
    { [ swap pick swap member? ] [ 3drop 0 ] }
    { [ member? ] [ 2 * ] }
    [ ]
  } cond ;

: (select-target) ( enemyarmy armygroup -- enemyarmy' from to )
  [ 
    over empty? [ f ] [
     2dup [ swap [ damage ] [ effective-power ]  [ initiative>> ] tri 3array ] curry supremum-by
     2dup damage zero? [ drop f ] when
    ] if
    [ nip swap remove ] keep
  ] keep swap ;

: select-target ( immunearmy infectionarmy armygroup -- immunearmy' infectionarmy' from to )
dup army>> "Immune System" = [ (select-target) ] [ swapd (select-target) [ swap ] 2dip ] if ;
 
: armygroup-key ( armygroup -- key  )
  [ army>> ] [ id>> number>string ] bi append ;
: select-targets ( army1 army2 -- targets )
2dup append [ [ effective-power ] [ initiative>> ] bi 2array ] inv-sort-with
[ select-target [ armygroup-key ] dip ] H{ } map>assoc 2nip ;

: remove-deads ( immunearmy infectionarmy dead-target -- immunearmy' infectionarmy' )
dup army>> "Immune System" = [ swapd swap remove swap ] [ swap remove ] if ;
: attack-target ( immunearmy infectionarmy targets armygroup -- immunearmy' infectionarmy' )
  swap over armygroup-key of [
   [ damage ] keep dup hp>> [ swapd /i - ] curry with change-units
   dup units>> 0 <= [ remove-deads ] [ drop ] if
  ] [ drop ] if* ;
: attack-targets ( army1 army2 targets -- army1' army2' )
2over append [ initiative>> ] inv-sort-with
[ attack-target ] with each ;

: day24-p1 ( -- n )
"/tmp/input" ascii file-lines
{ "" } split
[
 unclip but-last swap
 [ 1 + rot [ parse-nums ] [ parse-parens ] bi <armygroup> ]
 with map-index
] map first2 2dup . .
[ 2dup [ empty? ] bi@ or ] [
2dup select-targets attack-targets
] until
dup empty? [ drop ] [ nip ] if [ units>> ] map-sum ;

day24-p1 .