USING: accessors arrays assocs assocs.extras combinators.extras
combinators.short-circuit grouping kernel math math.functions
path-finding project-euler.common ranges sequences
sequences.extras ;
IN: project-euler.061

: n>x ( n s -- x )
    [ [ sq ] dip 2 - * ] [ 4 - * ] 2bi - 2 / ;

: x>n ( x s -- n )
    [ 2 - 8 * * ] [ 4 - sq + sqrt ] [ 4 - + ] [ 2 - 2 * / ] quad
    ;

: lower-n-limit ( s -- n ) 1000 swap x>n ceiling >integer ;

: upper-n-limit ( s -- n ) 9999 swap x>n floor >integer ;

:: 4-digit-s-gonal-numbers ( s -- set )
    s [ lower-n-limit ] [ upper-n-limit ] bi [a..b]
    [ s tuck n>x 2array ] map [ second 100 mod 10 >= ] filter ;

: number-sets ( -- sets )
    3 8 [a..b] [ 4-digit-s-gonal-numbers ] map-concat ;

SYMBOL: goal
TUPLE: cycle < astar sets ;
FORGET: cycle?
: cycle? ( node -- ? )
    { [ length 6 = ] [ 2 circular-clump [
        values first2 [ 100 mod ] [ 100 /i ] bi* =
    ] all? ] } 1&& ;
: <cycle> ( sets -- astar ) [ cycle new ] dip >>sets ;
M: cycle cost ( from to astar -- n ) 3drop 1 ;
M: cycle heuristic ( from to astar -- n ) 3drop 1 ;
M:: cycle neighbors ( node astar -- nodes )
    astar sets>> node keys '[ _ member? ] reject-keys
    node tuck last second 100 mod '[ 100 /i _ = ] filter-values
    [ suffix dup cycle? [ drop goal ] when ] with map ;

: euler061 ( -- n )
    number-sets dup [ 8 = ] filter-keys [
        1array goal rot <cycle> find-path
    ] with map-find drop but-last last values
    [ sum ] [ [ first 100 /i ] [ last 100 mod 100 * ] bi + + ]
    bi ;

SOLUTION: euler061