! Copyright (C) 2022 Doug Coleman.
! See http://factorcode.org/license.txt for BSD license.
USING: arrays combinators.extras grouping kernel math
math.parser sequences shuffle splitting vectors ;
IN: 2022-aoc5

: ntimes ( quot n -- ) swap times ; inline
: nths-of ( seq indices -- elts ) swap nths ; inline
: nth-of ( seq n -- elt ) swap nth ; inline

: parse-crate-line ( string -- seq )
    4 group harvest
    [ [ "[] " member? ] trim ] { } map-as ;

: parse-move-line ( string -- count,from,to )
    " " split { 3 5 1 } nths-of
    [ dec> ] map
    first3 [ [ 1 - ] bi@ ] dip 3array ;

: parse-aoc5-lines ( string -- crates insns )
    "\n\n" split1 [
        split-lines but-last [ parse-crate-line ] map flip
        [ reverse >vector harvest ] V{ } map-as
    ] [ split-lines [ parse-move-line ] map ] bi* ;

: do-move ( seqs from to n -- )
    [
        2array nths-of first2
        '[ _ _ [ pop ] [ push ] bi* ]
    ] dip ntimes ;

: do-move2 ( seqs from to n -- )
    [ 2array nths-of first2 V{ } clone ] dip
    [ nipd [ '[ _ pop _ push ] ] dip ntimes ]
    [ drop nipd reverse swap push-all ] 4bi ;

: aoc5-1 ( string -- seqs )
    parse-aoc5-lines
    [ dupd first3 do-move ] each
    [ last ] map concat ;

: aoc5-2 ( string -- seqs )
    parse-aoc5-lines
    [ dupd first3 do-move2 ] each
    [ last ] map concat ;