\ Forth Go - bitmaps
\ By Ian Osgood <iano@quirkster.com>

\ TODO: check some references:
\ Ting, C.H., Go in Forth. Dr. Dobb's Journal, No. 83, pp. 54 -65, September 1983.
\ [Ting 83] Go in Forth. Ting, C. H. Dr. Dobb's Journal, 83 (Sept. 1983), 54-60.
\  duplicated in Dr. Dobbs Toolbook of Forth (Ch. 4)  (only 8 screens of code!)
\  Bill Spight

\ IDEA: single board, two bits per square
\  00 empty
\  01 white
\  10 black
\  11 edge

\ IDEA: specify a smaller pattern's y origin and extent.
\   outside bounds, assumed to be zero
\   only valid for sources?

\ : bd-bounds ( b -- b limit start )
\   DUP @ IF dim 1+ 1
\   ELSE DUP @ 28 RSHIFT OVER @ 24 RSHIFT THEN ;

9 CONSTANT dim
dim 2 + CELLS CONSTANT bd-size
2 dim LSHIFT 2 - CONSTANT row-mask

: CELL-  [ 1 CELLS ] LITERAL - ;

\  board starts at 1 LSHIFT
\  edge is 1 dim 2 + LSHIFT 1 OR (0x401)
\  loop is dim 1+ 1 DO ... LOOP

\ board: array of bitmaps
\ location: y CELLS + @ x RSHIFT 1 AND

CREATE white bd-size ALLOT
CREATE black bd-size ALLOT
CREATE empty bd-size ALLOT

black VALUE tomove
white VALUE enemy

: switch-colors  tomove enemy TO tomove TO enemy ;
: .tomove
  tomove white = IF ." White(O)" ELSE ." Black(@)" THEN ."  to move" ;

\ xy 1..9   ( 0 and 10 are edges )
\ : I>xy   DUP 31 AND SWAP 5 RSHIFT ;
\ : bd@ ( I b -- ? )
\   OVER 5 RSHIFT CELLS + @
\   SWAP 31 AND RSHIFT 1 AND ;

: bd-set ( x y b -- )
  SWAP CELLS +  ( x r )
  1 ROT LSHIFT
  OVER @ OR SWAP ! ;
: bd-clr ( x y b -- )
  SWAP CELLS +  ( x r )
  1 ROT LSHIFT INVERT
  OVER @ AND SWAP ! ;
: bd-@ ( x y b -- 0/1 )
  SWAP CELLS + @ ( x r )
  SWAP RSHIFT 1 AND ;

: .board
  dim 1+ 1 DO CR
    dim 1+ 1 DO
      I J white bd-@ IF [CHAR] O EMIT THEN
      I J black bd-@ IF [CHAR] @ EMIT THEN
      I J empty bd-@ IF [CHAR] . EMIT THEN
      SPACE
    LOOP
  LOOP
  ." | " .tomove ;

\ board operations

: bd-copy ( src dest -- ) bd-size MOVE ;
: bd-erase ( dest -- ) bd-size ERASE ;

: bd-or ( src dest -- )
  dim 0 DO
    SWAP CELL+ SWAP CELL+
    OVER @ OVER @ OR OVER !
  LOOP 2DROP ;

: bd-and ( src dest -- )
  dim 0 DO
    SWAP CELL+ SWAP CELL+
    OVER @ OVER @ AND OVER !
  LOOP 2DROP ;

: bd-xor ( src dest -- )
  dim 0 DO
    SWAP CELL+ SWAP CELL+
    OVER @ OVER @ XOR OVER !
  LOOP 2DROP ;

: bd-not ( dest -- )
  dim 0 DO
    CELL+
    DUP @ INVERT row-mask AND OVER !
  LOOP DROP ;

: bd-2/ ( dest -- )
  dim 0 DO
    CELL+
    DUP @ 2/ row-mask AND OVER !
  LOOP DROP ;

: bd-2* ( dest -- )
  dim 0 DO
    CELL+
    DUP @ 2* row-mask AND OVER !
  LOOP DROP ;

: bd-r/ ( dest -- ) CELL+ DUP CELL+ SWAP dim CELLS MOVE ;
: bd-r* ( dest -- ) DUP CELL+ dim CELLS MOVE ;

: bd-expand ( dest -- )
  0 HERE ! HERE CELL+ OVER CELL+ ( dest tmp src )
  dim 0 DO
    DUP @ DUP 2/ SWAP 2* OR row-mask AND ( tmp src exp )
    OVER CELL- @ OR OVER CELL+ @ OR
    ROT TUCK !
    CELL+ SWAP CELL+
  LOOP DROP 0 SWAP !
  HERE SWAP bd-or ;

\ count and tests

: count-table CREATE ( bits -- )
  1 SWAP LSHIFT 0 DO
    0 I
    BEGIN DUP WHILE DUP 1- AND SWAP 1+ SWAP REPEAT
    DROP C,
  LOOP
  DOES> ( b -- count ) + C@ ;

dim 1+ count-table row-count

: bd-count ( b -- n ) 0
  dim 0 DO
    SWAP CELL+ DUP @ row-count ROT +
  LOOP NIP ;

: bd-empty? ( b -- ? )
  dim 0 DO
    CELL+
    DUP @ IF DROP FALSE UNLOOP EXIT THEN
  LOOP
  DROP TRUE ;

: bd-equals? ( b c -- ? )
  dim 0 DO
    CELL+ SWAP CELL+
    OVER @ OVER @ <> IF 2DROP FALSE UNLOOP EXIT THEN
  LOOP
  2DROP TRUE ;

: bd-intersects? ( b c -- ? )
  dim 0 DO
    CELL+ SWAP CELL+
    OVER @ OVER @ AND IF 2DROP TRUE UNLOOP EXIT THEN    
  LOOP
  2DROP FALSE ;

\ board stack  ( stack contents in brackets [])

CREATE bd-stack bd-size 8 * ALLOT
bd-stack VALUE bd-next
0 VALUE bd-top

: bnext ( [] -- [?] ) bd-next DUP TO bd-top bd-size + TO bd-next ;
: bdrop ( [b] -- [] ) bd-top DUP TO bd-next bd-size - TO bd-top ;

: blit ( b -- [b] ) bnext  bd-top bd-copy ;
: bnip ( [a b] -- [b] ) bdrop  bd-next bd-top bd-copy ;
: bdup ( [b] -- [b b] ) bd-top blit ;
: bover ( [a b] -- [a b a] ) bd-top bd-size - blit ;
\ : bswap ( -- ) ;

: bor  ( [a b] -- [a|b] ) bdrop bd-next bd-top bd-or ;
: bxor ( [a b] -- [a^b] ) bdrop bd-next bd-top bd-xor ;
: band ( [a b] -- [a&b] ) bdrop bd-next bd-top bd-and ;

: bover-and ( [a b] -- [a a&b] ) bd-top bd-size - bd-top bd-and ;

: bnot ( [b] -- [b] ) bd-top bd-not ;
: b2* ( [b] -- [b] ) bd-top bd-2* ;
: b2/ ( [b] -- [b] ) bd-top bd-2/ ;
: br* ( [b] -- [b] ) bd-top bd-r* ;
: br/ ( [b] -- [b] ) bd-top bd-r/ ;

\ : b= ( [a b] -- ? ) bdrop bd-next bd-top bd-equals? bdrop ;
: b0= ( [b] -- ? ) bdrop bd-next bd-empty? ;
: bcount ( [b] -- n ) bdrop bd-next bd-count ;

: b0 ( -- [0] ) bnext  bd-top bd-erase ;
: bbit ( x y -- [b] ) b0  bd-top bd-set ;

\ derived boards

\ expand: basis of string connection, liberties, capture etc.

\  0 0 0      0 1 0
\  0 1 0  ->  1 1 1
\  0 0 0      0 1 0

: expand ( [b] -- [e] ) bd-top bd-expand ;
\  bdup b2*  bover b2/ bor  bover br* bor  bover br/ bor  bor ;

\ isolate string

\ b w   0 0    0 0    0 1
\ w w : 1 0 -> 1 1 -> 1 1
\ w .   0 0    1 0    1 0

: string ( [b] x y -- [s] count )	\ isolate the string at x y
  \ 2DUP bd-top bd-@ 0= IF 2DROP 0 EXIT THEN
  bbit bover-and 1 ( [b s] bitcount )
  BEGIN
    expand bover-and
    bd-top bd-count TUCK =
  UNTIL bnip ;

\ liberties of a string
\ liberties == neighbors & empty

\ b w .    0 0 1
\ . w . -> 1 0 1
\ . b b    0 0 0

: liberties ( [s] -- [l] )
  expand empty bd-top bd-and ;

: capture ( s -- )
  DUP empty bd-or  enemy bd-xor ;

\ check for liberties at each dilation for an early cutoff
: ?capture ( x y -- )
  2DUP enemy bd-@ IF
    bbit 1 ( [s] count )
    BEGIN
      expand
      empty bd-top bd-intersects? IF		\ liberties? exit
        bdrop DROP EXIT THEN
      enemy bd-top bd-and
      bd-top bd-count TUCK =
    UNTIL					\ no liberties: capture
    bd-top capture bdrop DROP
  ELSE
    2DROP THEN ;

: check-captures ( x y -- )
  OVER 1+ OVER ?capture
  OVER 1- OVER ?capture
  2DUP 1+ ?capture
       1- ?capture ;

\ . . . .
\ . O O @
\ O . O @  <
\ . O . @

\ find the territory of a point

: eyespace ( x y -- [e] count )      \ isolate the space at x y
  \ 2DUP bd-top bd-@ 0= IF 2DROP 0 EXIT THEN
  bbit empty bd-top bd-and 1 ( [e] bitcount )
  BEGIN
    liberties
    bd-top bd-count TUCK = OVER 9 > OR
  UNTIL ;

: is-eye? ( x y -- ? )
  eyespace
  DUP 0= SWAP 9 > OR IF bdrop FALSE EXIT THEN
  expand enemy bd-top bd-and b0= ;

\ O @ @ .
\ . O @ @  <
\ O . O @  <
\ . O @ @

\ rule: false eye if it is the only liberty of any neighbor

: atari? ( x y -- ? )
  2DUP tomove bd-@ IF
    tomove blit string DROP
    liberties bcount 1 =
  ELSE 2DROP FALSE THEN ;

: is-false-eye? ( x y -- ? )
  OVER 1+ OVER atari? IF 2DROP TRUE EXIT THEN
  OVER 1- OVER atari? IF 2DROP TRUE EXIT THEN
  2DUP 1+      atari? IF 2DROP TRUE EXIT THEN
       1-      atari? ;

\ simple scoring (no seki, no ko, all groups settled)

: score-tomove ( -- [t] n )
  tomove blit bd-top bd-count
  BEGIN
    expand enemy blit bnot band
    bd-top bd-count TUCK =
  UNTIL ;

: scores ( -- b w )
  score-tomove bdrop switch-colors
  score-tomove bdrop switch-colors
  tomove white = IF SWAP THEN ;

6 VALUE komi		\ score adjustment for black having the first move

: score ( -- n ) scores - komi - ;	\ score relative to black

: .scores ( -- ) scores
  ." White(O): " . ." +" komi . ." komi" ."   Black(@): " . ;

\ history stack (black and white only)
\ TODO: overflow checking
bd-size 2* CONSTANT hist-size
CREATE bd-hist dim dim * 3 * hist-size * ALLOT
bd-hist VALUE bd-hist-top

: hist-push ( b -- )
  bd-hist-top bd-copy  bd-hist-top bd-size + TO bd-hist-top ;
: hist-pop ( b -- )
  bd-hist-top bd-size - TO bd-hist-top  bd-hist-top SWAP bd-copy ;
: hist-save ( -- )
  white hist-push black hist-push ;
: hist-restore ( -- )
  bd-hist-top bd-hist = IF EXIT THEN
  black hist-pop white hist-pop
  black empty bd-copy  white empty bd-or  empty bd-not ;

: reset
  white bd-erase black bd-erase empty DUP bd-erase bd-not
  bd-hist TO bd-hist-top
  black TO tomove  white TO enemy ;

: hist= ( -n -- ? )
  hist-size * bd-hist-top +
  DUP       white bd-equals? 0= IF DROP FALSE EXIT THEN
  bd-size + black bd-equals? ;

\ simple test: did last position repeat?
: is-ko? ( -- ? )
  bd-hist-top bd-hist - hist-size 2* < IF FALSE EXIT THEN
  -2 hist= ;

\ two passes mean the game is over
: game-over? ( -- ? )
  bd-hist-top bd-hist - hist-size 2* < IF FALSE EXIT THEN
  -1 hist= IF -2 hist= ELSE FALSE THEN ;

\ quicker test used in the playouts
VARIABLE passes
: done?
  passes @ 1 > ;

: pass ( -- )
  1 passes +! hist-save switch-colors ;

: makemove ( x y -- )
  2DUP D0= IF 2DROP pass EXIT THEN
  hist-save
  2DUP tomove bd-set 2DUP empty bd-clr
  2DUP check-captures
  switch-colors
  ?capture ;		\ check for suicide

: takeback ( -- )
  hist-restore switch-colors ;

: move-ok? ( x y -- ? )
  2DUP empty bd-@ 0= IF 2DROP FALSE EXIT THEN
  2DUP is-eye? IF
    2DUP is-false-eye? 0= IF 2DROP FALSE EXIT THEN
  THEN
  2DUP makemove
  empty bd-@ IF takeback FALSE EXIT THEN
  is-ko? IF takeback FALSE EXIT THEN
  TRUE ;

\ setup takes pasted board using accept
: setup
  reset  empty bd-erase
  dim 1+ 1 DO
    PAD 80 ACCEPT  dim 2* < ABORT" Bad board format!"
    dim 1+ 1 DO
      PAD I 1- 2* + C@
      DUP [CHAR] @ = IF DROP I J black bd-set ELSE
      DUP [CHAR] O = IF DROP I J white bd-set ELSE
      DUP [CHAR] . = IF DROP I J empty bd-set ELSE
        [CHAR] ' EMIT EMIT [CHAR] ' EMIT -1 ABORT" Bad board character!"
      THEN THEN THEN
    LOOP
  LOOP \ check for color to move on last line
  PAD dim 2* + 8 + C@ [CHAR] O = IF switch-colors THEN
  .board ;

\ random
HERE VALUE seed
: RANDOM ( -- u ) seed $107465 * $234567 + DUP TO seed ;
: choose ( n -- 0..n-1 ) RANDOM UM* NIP ;

CREATE candidates bd-size ALLOT
: nth-candidate
  candidates
  dim 1+ 1 DO
    CELL+ TUCK @ row-count -
    DUP 0< IF
      1+ SWAP @
      dim 1+ 1 DO
        2/ DUP 1 AND IF
          OVER 0= IF
            2DROP I J UNLOOP UNLOOP EXIT
          THEN
          SWAP 1+ SWAP
        THEN
      LOOP
    THEN
    SWAP
  LOOP  ." nth-candidate failed!" 2DROP 0. ;

: random-candidate          \ side-effect: makes a move
  empty candidates bd-copy
  candidates bd-count
  BEGIN
    BEGIN
      DUP choose nth-candidate ( count x y )
      \ CR .s
      2DUP is-eye?
    WHILE
      \ ." eye"
      2DUP is-false-eye? IF ( ." [false]") FALSE ELSE TRUE THEN
    WHILE
      candidates bd-clr 1-
      DUP 0= IF ( ."  PASS") pass DROP EXIT THEN
    REPEAT THEN
    2DUP makemove
    2DUP empty bd-@  DUP IF ( ." suicide")
    ELSE DROP is-ko? DUP IF ( ." ko") THEN
    THEN
  WHILE
    takeback
    candidates bd-clr 1-
    DUP 0= IF ( ."  PASS") pass DROP EXIT THEN
  REPEAT 2DROP DROP
  0 passes ! ;

: playout ( -- score )
  bd-hist-top enemy tomove
    BEGIN random-candidate done? UNTIL score >R
  TO tomove TO enemy
  hist-size + TO bd-hist-top hist-restore
  R> ;

\ plain Monte Carlo (best of avg N playouts)

: playouts ( n -- sum-score )
  0 SWAP 0 DO playout + LOOP ;

\ try moves from the center out to the edges
HEX
: coordinate-table CREATE DOES> ( I -- x y ) SWAP CELLS + @ 10 /MOD ;
coordinate-table central
  55 , 45 , 54 , 65 , 56 , 46 , 44 , 64 , 66 ,
  35 , 34 , 33 , 43 , 53 , 63 , 73 , 74 , 75 ,
  76 , 77 , 67 , 57 , 47 , 37 , 36 ,  25 , 24 ,
  23 , 22 , 32 , 42 , 52 , 62 , 72 , 82 , 83 ,
  84 , 85 , 86 , 87 , 88 , 78 , 68 , 58 , 48 ,
  38 , 28 , 27 , 26 ,  51 , 41 , 31 , 21 , 12 ,
  13 , 14 , 15 , 16 , 17 , 18 , 92 , 93 , 94 ,
  95 , 96 , 97 , 98 , 89 , 79 , 69 , 59 , 49 ,
  39 , 29 , 81 , 71 , 61 ,  11 , 91 , 19 , 99 ,
DECIMAL

4 VALUE mc-limit
: best-playout ( -- x y )
  0 0 dim dim * mc-limit * NEGATE ( best move and score )
  dim dim * 0 DO
    I central move-ok? IF
      mc-limit playouts takeback
      tomove white = IF NEGATE THEN
      2DUP < IF
        NIP NIP NIP  I central ROT
        CR .s
        DUP dim dim * mc-limit * = IF LEAVE THEN
      ELSE DROP THEN
    THEN
  LOOP CR ." Best = " . ." : " OVER . DUP . ;

\ UCT search

: U>F ( n -- f ) 0 D>F ;

\ node layout: siblings, children, visits, wins, move x,y

: n>next              @ ;
: n>child    CELL+    @ ;
: n>visits 2 CELLS +  @ ;
: n>wins   3 CELLS +  @ ;
: n>move   4 CELLS + 2@ ;

: n>winrate  DUP n>wins U>F n>visits U>F f/ ;

: n>children! CELL+   ! ;
: n>visited! 2 CELLS +  1 SWAP +! ;
: n>won!     3 CELLS +  1 SWAP +! ;

: root-node ( -- node )
  HERE 0 , 0 , 1 , 0 , 0. 2, ;	\ one visit, so we don't playout at root

\ expand children if needed

: ?children ( node -- node )
  DUP n>child 0= IF
    0				( tail )
    empty candidates bd-copy
    candidates bd-count         ( tail count )
    BEGIN DUP WHILE	\ for all candidates
      DUP choose nth-candidate  ( tail count x y )
      2DUP move-ok? IF
        takeback
        2SWAP HERE SWAP ROT , 0 , 0 , 0 , 2SWAP 2DUP 2,
      THEN
      candidates bd-clr 1-
    REPEAT DROP
    DUP 0= IF
      DROP HERE 0 , 0 , 0 , 0 , 0. ( pass ) 2,
    THEN
    OVER n>children!
  THEN ;

: .node ( node -- )
  ." (winrate:" DUP n>wins 1 U.R ." /" DUP n>visits U.
  ." mv: "  n>move SWAP . . ." ) " ;
: .children ( node -- )
  CR n>child BEGIN DUP WHILE DUP .node n>next REPEAT DROP ;

FVARIABLE best
VARIABLE best-node
: pick-best-move ( root -- x y )
  0e best F!
  n>child
  BEGIN DUP WHILE
    DUP n>winrate
    best F@ FOVER F< IF
      best F!
      DUP best-node !
    ELSE FDROP THEN
    n>next
  REPEAT DROP
  best-node @ CR ." Best move: " .node
  best-node @ n>move ;

\ Two states:
\ 1. just expanded, visit each child once in random order
\    achieved by ?children building the child list in random order

: all-select ( parent -- child )
  n>child
  BEGIN DUP n>next ( c n )
        DUP
  WHILE DUP n>visits 0=
  WHILE NIP
  REPEAT THEN DROP ;

\  2. all children visited at least once, use the UCT metric

: uct-select ( parent -- child )
  DUP n>visits U>F FLOG  ( F: log(parent visits)
  0e best F!
  n>child
  BEGIN DUP WHILE
    FDUP ( log[parent visits] ) DUP n>visits 5 * U>F F/ FSQRT   ( F: uct )
    \ SELECTIVE-FACTOR F*   \ <1 selective, >1 uniform
    DUP n>winrate F+     ( F: VALUE )
    best F@ FOVER F< IF
      best F!
      DUP best-node !
    ELSE FDROP THEN
    n>next
  REPEAT DROP FDROP
\  DEPTH 2 > IF
\    best-node @ CR ." Expanding(" DEPTH 2 - . ." ): " .node THEN
  best-node @ ;

: move-select ( parent -- child )
  ?children
  DUP n>child n>visits IF uct-select ELSE all-select THEN ;

: uct-playout ( node -- result )
  game-over? IF
    score tomove white = IF 0> ELSE 0< THEN
  ELSE DUP n>visits IF
    DUP move-select ( node next )
    DUP n>move makemove RECURSE takeback
  ELSE
    playout tomove white = IF 0> ELSE 0< THEN	\ convert to win/loss
  THEN THEN ( node result )
  DUP IF OVER n>won! THEN
  SWAP n>visited!
  0= ;				\ colors alternate, so swap result for parent node

   10 VALUE    time-limit
10000 VALUE playout-limit

: seconds ( -- n ) TIME&DATE 2DROP 1- 24 * + 60 * + 60 * + ;

: uct-search ( -- x y )
  root-node seconds time-limit + 0
  BEGIN
    ROT DUP uct-playout DROP -ROT
    1+ DUP playout-limit <
  WHILE
    OVER seconds >
  WHILE REPEAT THEN
  . ." playouts in " seconds SWAP - time-limit + . ." seconds"
  DUP pick-best-move
  ROT HERE - DUP NEGATE . ." bytes used" ALLOT ;	\ assert HERE is restored

\ commands
\  also .board .tomove .scores

: rm ( -- ) random-candidate .board ;			\ random move

\ configure with N TO mc-limit
: mc ( -- ) best-playout makemove .board ;		\ Monte Carlo

\ configure with N TO playout-limit and S TO time-limit
: uct ( -- ) \ seconds begin dup seconds <> until drop
  uct-search makemove .board ;

DEFER go  ' uct IS go	\ go player can be rm, mc, or uct

\ make a move
\  x,y = [1..dim] where 1 1 is top left of .board
: mv ( x y -- )
  move-ok? IF .board CR go ELSE ." Bad move!" THEN ;

: self-play ( p1 p2 -- )
  BEGIN SWAP DUP EXECUTE game-over? UNTIL 2DROP CR .scores ;

\ init

reset .board CR