USING: accessors fry kernel locals make math multiline
prettyprint sequences shuffle slots.syntax ;
QUALIFIED: math
! QUALIFIED-WITH: math.order mo
IN: binary-tree

TUPLE: node parent left right { size-below initial: 0 } value ;
TUPLE: binary-tree { root node } ;


! : <=> ( node1 node2 -- ? )
!    [ value>> ] bi@ mo:<=> ;
! : > ( node1 node2 -- ? ) 
!    [ value>> ] bi@ math:> ; 

! : = ( node1 node2 -- ? )
!     [ value>> ] bi@ kernel:= ;

: < ( node1 node2 -- ? ) 
    [ value>> ] bi@ math:< ;
    
: >= ( node1 node2 -- ? ) 
    < not ;
! : <= ( node1 node2 -- ? ) 
!    > not ;


: <node> ( value -- node )    
    node new 
    swap >>value ;
    
: <binary-tree> ( node -- btree )
    binary-tree new swap >>root ;

    
<PRIVATE

: inc-parent-size ( child -- )
    parent>> dup 
    [
        [ 1 + ] change-size-below inc-parent-size 
    ] [ drop ] if ;
    
: add-child-on-right ( node child -- )
    over >>parent 
        [ >>right drop ]
        [ nip inc-parent-size ] 2bi ;
: add-child-on-left ( node child -- )
    over >>parent 
        [ >>left drop ]
        [ nip inc-parent-size ] 2bi ;
PRIVATE>

GENERIC: add-node ( child node -- )


M: node add-node ( child node -- )
    2dup >= [ ! go right
        dup right>> ! can be node/f
            [ right>> add-node ] [ swap add-child-on-right ] if
    ] [ ! go left
        dup left>> 
            [ left>> add-node ] [ swap add-child-on-left ] if
    ] if ;

M: binary-tree add-node ( child node -- )
    root>> add-node ;
    
: add-value ( node value -- )
    <node> swap add-node ;
    
GENERIC: min ( node -- min-node )
GENERIC: max ( node -- max-node )

M: node min ( node -- min-node )
    dup left>> [ left>> min ] when ;

M: node max ( node -- max-node )
    dup right>> [ right>> max ] when ;

M: binary-tree min ( node -- min-node )
    root>> min ;
    
M: binary-tree max ( node -- max-node )
    root>> max ;

: root ( node -- root )
    dup parent>> [ parent>> root ] when ;
    
<PRIVATE
: (get-list) ( node -- )
    [ left>> dup [ (get-list) ] [ drop ] if ]
    [ , ] 
    [ right>> dup [ (get-list) ] [ drop ] if ] tri ;
PRIVATE>   

GENERIC: items ( base-node -- items  )
M: node items ( base-node -- seq )
    [ (get-list) ] { } make ;
M: binary-tree items ( base-node -- seq )    
    root>> items ;

: values ( btree -- seq )
    items [ value>> ] map ;
    
: print ( binary-tree -- )
    values . ;
    
GENERIC: find-node* ( value node -- node/f )    

: find-node ( value node -- node/f )    
    dup [ find-node* ] [ 2drop f ] if ;
    
M:: node find-node* ( value node -- node/f )    
    node value>> value = 
    [ node ]
    [ value node left>> find-node :> left-result
        left-result dup [ drop value node right>> find-node :> right-result 
        right-result dup [ drop f ] unless
        ] unless
    ] if ;
    
M: binary-tree find-node* ( value btree -- node/f )
    root>> find-node ;    

    
: add-values ( node seq -- node' )
    dupd [ add-value ] with each ;

: add-nodes ( items node -- node' )
    tuck '[ _ add-node ] each ;
    
: children ( node -- seq )
    slots{ right left } ;

: right? ( child -- ? )
    dup parent>> right>> eq? ;

<PRIVATE
: rm-child ( node -- )
    [ parent>> ] [ right? ] bi 
    [ f >>right ] [ f >>left ] if drop ;
PRIVATE>

: remove ( node -- )
    dup [
    [ children harvest ] [ parent>> ] [ rm-child ] tri add-nodes 
    ] when drop ;

: remove-nodes ( nodes -- )
    [ remove ] each ;
    
: find-nodes ( btree values -- nodes )
    swap [ find-node ] curry map ;
    
: remove-values ( btree values -- )
    find-nodes remove-nodes ;