USING: arrays assocs assocs.extras io.encodings.ascii io.files
kernel make math.order math.parser namespaces prettyprint
sequences sets splitting ;
FROM: namespaces => set ;
IN: aoc.2020.16

SYMBOLS: fields ranges tickets ;        ! This problem fried my brain.
                                        ! I'm sure it's a bad way to solve it, but it works
: parse-input ( -- )                    ! and that's all that matters to me right now
    "input.txt" ascii file-lines
    { "" "nearby tickets:" } without
    { "your ticket:" } split1
    [
        [
            ": " split1 " " split1 " " split1 nip
            [ "-" split1 [ dec> ] bi@ 2array ] bi@ 2array
        ] H{ } map>assoc dup values concat ranges set fields set
    ]
    [ harvest [ "," split [ dec> ] map ] map tickets set ] bi* ;

: valid-field? ( n -- ? )
    ranges get [ swapd between? ] with assoc-any? ;

: part1 ( -- )
    parse-input tickets get concat [ valid-field? ] reject sum . ;

: set-valid-tickets ( -- )
    tickets [ [ [ valid-field? ] all? ] filter ] change ;

: valid-field?* ( assoc n -- ? )
    [ -rot between? ] curry assoc-any? ;

: valid-row? ( assoc seq -- ? ) [ valid-field?* ] with all? ;

: id-field ( seq -- key )
    fields get swap [ nipd valid-row? ] curry assoc-find 2drop
    dup fields get delete-at ;

: comb ( seq -- newseq )
    dup [ assoc-size 1 = ] find keys first tuck ,, swap
    [ [ delete-at ] keep ] with map ;

: id-fields ( seq -- assoc )
    [ [ dup [ assoc-empty? ] all? ] [ comb ] until ] H{ } make
    nip ;

: part2 ( -- )
    set-valid-tickets tickets get flip
    [ fields get swap [ valid-row? nip ] curry assoc-filter ]
    map id-fields [ "departure" head? ] filter-keys values
    tickets get first nths product . ;