! Tom Weeks - Math 2051, Section 01
! u06a2 (2009.11.15) - Programming Project 2 (Page 400)
!
! For a given transportation network, determine a maximal flow
! by repeated use of the flow augmentation algorithm, starting
! with the flow that is zero along every arc.

USING: accessors arrays combinators combinators.short-circuit
       io kernel math math.order math.parser prettyprint
       sequences sorting strings ;
IN: u06a2

TUPLE: arc { from-to   pair    initial: { "" "" } }
           { capacity  real                       }
           { flow      real                       } ;

TUPLE: label { predecessor string  }
             { +/-         string  }
             { flow-delta  real    } ;

TUPLE: vertex { name     string  }
              { label    label   }
              { scanned? boolean }
              { labeled# integer } ; ! 0 means unlabeled

TUPLE: network { vertices array }
               { arcs     array } ;
 
: vertex-names ( network -- v-names )
  vertices>> [ name>> ] map ;

: vertex-pairs ( network -- v-pairs )
  arcs>> [ from-to>> ] map ;

: predecessors ( network -- v-names )
  vertex-pairs [ first ] map ;

: successors ( network -- v-names )
  vertex-pairs [ second ] map ;

: this-vertex ( v-name network -- v ) 
  vertices>> [ name>> over = ] find 2nip ;

: source-name ( network -- v-name )
  [ successors ] [ vertex-names ] bi
  [ over member? not ] find 2nip ;

: sink-name ( network -- v-name )
  [ predecessors ] [ vertex-names ] bi
  [ over member? not ] find 2nip ;

: source ( network -- v ) [ source-name ] keep this-vertex ;

: sink ( network -- v ) [ sink-name ] keep this-vertex ;

: infinity ( -- n ) 2 26 shift dup 1 - + ; inline ! fixnum max

: <label> ( predecessor +/- flow-delta -- label ) label boa ;

: label-source ( network -- network' )
  dup source "source" "+" infinity <label> >>label
  1 swap (>>labeled#) ;

: labeled? ( v -- ? ) labeled#>> 0 > ;

: scanned? ( v -- ? ) scanned?>> ;

: labeled ( network -- vs ) vertices>> [ labeled? ] filter ;

: unlabeled ( network -- vs )
  vertices>> [ labeled? not ] filter ;

: highest-labeled# ( network -- n )
  vertices>> [ labeled#>> ] map natural-sort last ;

: next-label# ( network -- n ) highest-labeled# 1 + ;

: >name-pair ( v w -- v-pair ) [ name>> ] bi@ 2array ;
    
: is-arc? ( network v w -- ? )
  >name-pair swap arcs>> [ from-to>> ] map member? ;

: this-arc ( network v w -- arc )
  >name-pair swap arcs>> [ from-to>> over = ] filter first
  nip ;

: capacity-remaining ( arc -- n )
  [ capacity>> ] [ flow>> ] bi - ;

: capacity-remaining? ( arc -- ? ) capacity-remaining 0 > ; 

: is-arc-and-capacity-remaining? ( network v w -- ? )
  { [ is-arc? ] [ this-arc capacity-remaining? ] } 3&& ;

: is-arc-and-flow>0? ( network v w -- ? )
  { [ is-arc? ] [ this-arc flow>> 0 > ] } 3&& ;

: min(v-flow-delta,w-capacity-remaining) ( network v w -- n )
  [ this-arc capacity-remaining ] 2keep
  drop label>> flow-delta>>
  min ;

: min(v-flow-delta,w-to-v-flow) ( network v w -- n )
  [ swap this-arc flow>> ] 2keep
  drop label>> flow-delta>>
  min ;

: maybe-label-vertex ( network v w -- )
  { { [ 3dup is-arc-and-capacity-remaining? ]
      [ 3dup min(v-flow-delta,w-capacity-remaining)
        [ over name>> "+" ] dip <label> >>label
        pick next-label# >>labeled# ] }
    { [ 3dup swap is-arc-and-flow>0? ]
      [ 3dup min(v-flow-delta,w-to-v-flow)
        [ over name>> "-" ] dip <label> >>label
        pick next-label# >>labeled# ] }
    [ ]
  } cond 3drop ; 

: label-some-neighbors ( network v unlabeled -- network' )
  dup empty?
  [ drop t >>scanned? drop ]
  [ 3dup first maybe-label-vertex rest label-some-neighbors ]
  if ;

: first-labeled ( vs -- v ) [ labeled#>> ] sort-with first ;

: select-vertex-to-scan ( network -- v )
  vertices>>
  [ [ labeled? ] [ scanned? not ] bi and ] filter
  first-labeled ;

: scan-again? ( network -- ? )
  [ sink labeled? ]
  [ labeled [ scanned? ] all? ] bi or not ;

: scan-and-label ( network -- network' )
  dup [ select-vertex-to-scan ] [ unlabeled ] bi
  label-some-neighbors
  dup scan-again? [ scan-and-label ] when ;

: adjust-flow ( network v -- network' v' )
  dup label>> dup +/->> "+" =
  [ predecessor>> pick this-vertex 3dup swap this-arc
    [ nip ] dip pick sink label>> flow-delta>>
    over flow>> + >>flow ]
  [ predecessor>> pick this-vertex 3dup this-arc
    [ nip ] dip pick sink label>> flow-delta>>
    over flow>> swap - >>flow ]
  if drop over source over = [ adjust-flow ] unless ;

: increase-flow ( network -- network'/f )
  dup sink labeled? not
  [ drop f ] ! This flow is already a maximal flow.
  [ dup sink adjust-flow drop ]
  if ;

: augment-flow ( network -- network'/f )
  label-source scan-and-label increase-flow ; 

: delete-labels ( network -- network' )
  dup vertices>>
  [ T{ label } >>label
    f >>scanned?
    0 >>labeled# ]
  map drop ;

: print-vertex ( v -- )
  dup name>> ": " append write
  label>>
  "(" write
  [ predecessor>> "," append ]
  [ +/->> "," append ]
  [ flow-delta>> number>string ")" append ]
  tri [ write ] tri@ nl ;

: print-vertices ( network -- )
  "Vertices: " print vertices>> [ print-vertex ] each ;

: print-arc ( arc -- )
  dup from-to>>
  [ first "," append write ] [ second ": " append write ] bi
  [ capacity>> number>string "," append write ]
  [ flow>> number>string print ]
  bi ;

: print-arcs ( network -- )
  "Arcs: " print arcs>> [ print-arc ] each ;

: print-network ( network iteration# -- )
  "============" print
  [ "Iteration " ] dip number>string ": " [ write ]
  tri@ nl "============" print nl
  [ print-vertices nl ] [ print-arcs ] bi nl nl ;

: maximize-flow' ( network iteration# -- network' iteration# )
  1 + over augment-flow
  [ 2dup print-network
    [ delete-labels ] dip maximize-flow' ]
  when ;

: maximize-flow ( network -- ) 0 maximize-flow' 2drop ;

: fig7.12 ( -- network ) ! p.372
  T{ network f { T{ vertex f "A" }
                 T{ vertex f "B" }
                 T{ vertex f "C" }
                 T{ vertex f "D" }
                 T{ vertex f "E" } }
               { T{ arc f { "A" "B" } 6 }
                 T{ arc f { "A" "C" } 8 }
                 T{ arc f { "A" "D" } 3 }
                 T{ arc f { "B" "C" } 5 }
                 T{ arc f { "B" "D" } 6 }
                 T{ arc f { "C" "E" } 4 }
                 T{ arc f { "D" "E" } 9 } } } ;

: fig7.21 ( -- network ) ! p.375
  T{ network f { T{ vertex f "A" }
                 T{ vertex f "B" }
                 T{ vertex f "C" }
                 T{ vertex f "D" }
                 T{ vertex f "E" }
                 T{ vertex f "F" }
                 T{ vertex f "G" } }
               { T{ arc f { "A" "B" } 5 }
                 T{ arc f { "A" "C" } 7 }
                 T{ arc f { "A" "D" } 4 }
                 T{ arc f { "B" "C" } 1 }
                 T{ arc f { "B" "F" } 3 }
                 T{ arc f { "C" "D" } 2 }
                 T{ arc f { "C" "E" } 5 }
                 T{ arc f { "C" "F" } 4 }
                 T{ arc f { "D" "E" } 4 }
                 T{ arc f { "E" "F" } 1 }
                 T{ arc f { "E" "G" } 6 }
                 T{ arc f { "F" "G" } 9 } } } ;

: 7.2problem2 ( -- network ) ! p.379
  T{ network f { T{ vertex f "A" }
                 T{ vertex f "B" }
                 T{ vertex f "C" }
                 T{ vertex f "D" }
                 T{ vertex f "E" } }
               { T{ arc f { "A" "B" } 6.8 4.2 }
                 T{ arc f { "A" "C" } 2.2 1.8 }
                 T{ arc f { "A" "D" } 3.1 1.5 }
                 T{ arc f { "B" "C" } 5.2 3.3 }
                 T{ arc f { "B" "D" } 1.5 0.9 }
                 T{ arc f { "C" "E" } 7.6 5.1 }
                 T{ arc f { "D" "E" } 3.2 2.4 } } } ;

: 7.2problem8 ( -- network ) ! p.379
  T{ network f
    { T{ vertex f "A" T{ label f "B" "-" 1 } t 2 }
      T{ vertex f "B" T{ label f "D" "+" 4 } t 3 }
      T{ vertex f "C" T{ label f "B" "-" 3 } t 4 }
      T{ vertex f "D" T{ label f "source" "+" 268435455 } t 1 }
      T{ vertex f "E" T{ label f "C" "-" 3 } t 5 }
      T{ vertex f "F" T{ label f "E" "+" 3 } f 6 } }
    { T{ arc f { "A" "B" } 2 1 }
      T{ arc f { "B" "E" } 7 7 }
      T{ arc f { "C" "A" } 3 1 }
      T{ arc f { "C" "B" } 5 3 }
      T{ arc f { "D" "B" } 7 3 }
      T{ arc f { "D" "E" } 1 1 }
      T{ arc f { "D" "F" } 6 6 }
      T{ arc f { "E" "C" } 4 4 }
      T{ arc f { "E" "F" } 8 4 } } } ;

: 7.2problem11 ( -- network ) ! p.380
  T{ network f
    { T{ vertex f "A" }
      T{ vertex f "B" }
      T{ vertex f "C" }
      T{ vertex f "D" }
      T{ vertex f "E" }
      T{ vertex f "F" } }
    { T{ arc f { "A" "B" } 10 8 }
      T{ arc f { "A" "C" } 6 4 }
      T{ arc f { "B" "D" } 7 5 }
      T{ arc f { "B" "E" } 4 4 }
      T{ arc f { "C" "B" } 3 1 }
      T{ arc f { "C" "E" } 3 3 }
      T{ arc f { "D" "E" } 2 2 }
      T{ arc f { "D" "F" } 3 3 }
      T{ arc f { "E" "F" } 11 9 } } } ;

: 7.3problem8 ( -- network ) ! p.386
  T{ network f
    { T{ vertex f "A" T{ label f } }
      T{ vertex f "B" T{ label f } }
      T{ vertex f "C" T{ label f } }
      T{ vertex f "D" T{ label f } }
      T{ vertex f "E" T{ label f } }
      T{ vertex f "F" T{ label f } } }
    { T{ arc f { "A" "B" } 3 3 }
      T{ arc f { "A" "C" } 4 4 }
      T{ arc f { "A" "D" } 5 2 }
      T{ arc f { "B" "E" } 7 3 }
      T{ arc f { "D" "B" } 2 0 }
      T{ arc f { "D" "F" } 2 2 }
      T{ arc f { "E" "C" } 2 2 }
      T{ arc f { "E" "F" } 4 1 }
      T{ arc f { "F" "C" } 3 3 } } } ;

: 7.4problem16 ( -- network ) ! p.395
  T{ network f { T{ vertex f "C" }
                 T{ vertex f "D" }
                 T{ vertex f "E" }
                 T{ vertex f "F" }
                 T{ vertex f "G" }
                 T{ vertex f "S" }
                 T{ vertex f "T" }
                 T{ vertex f "U" }
                 T{ vertex f "V" }
                 T{ vertex f "W" }
                 T{ vertex f "s" }
                 T{ vertex f "t" } }
               { T{ arc f { "s" "S" } 1 }
                 T{ arc f { "s" "T" } 1 }
                 T{ arc f { "s" "U" } 1 }
                 T{ arc f { "s" "V" } 1 }
                 T{ arc f { "s" "W" } 1 }
                 T{ arc f { "C" "t" } 1 }
                 T{ arc f { "D" "t" } 1 }
                 T{ arc f { "E" "t" } 1 }
                 T{ arc f { "F" "t" } 1 }
                 T{ arc f { "G" "t" } 1 }
                 T{ arc f { "S" "E" } 1 }
                 T{ arc f { "S" "F" } 1 }
                 T{ arc f { "T" "C" } 1 }
                 T{ arc f { "T" "D" } 1 }
                 T{ arc f { "T" "G" } 1 }
                 T{ arc f { "U" "E" } 1 }
                 T{ arc f { "U" "F" } 1 }
                 T{ arc f { "V" "C" } 1 }
                 T{ arc f { "V" "D" } 1 }
                 T{ arc f { "V" "F" } 1 }
                 T{ arc f { "V" "G" } 1 }
                 T{ arc f { "W" "C" } 1 }
                 T{ arc f { "W" "E" } 1 }
                 T{ arc f { "W" "F" } 1 } } } ;

: 7.4problem22 ( -- network ) ! p.395
  T{ network f { T{ vertex f "A" }
                 T{ vertex f "B" }
                 T{ vertex f "C" }
                 T{ vertex f "D" }
                 T{ vertex f "E" }
                 T{ vertex f "F" }
                 T{ vertex f "G" }
                 T{ vertex f "H" }
                 T{ vertex f "I" }
                 T{ vertex f "J" }
                 T{ vertex f "s" }
                 T{ vertex f "t" } }
               { T{ arc f { "s" "A" } 1 }
                 T{ arc f { "s" "B" } 1 }
                 T{ arc f { "s" "C" } 1 }
                 T{ arc f { "s" "D" } 1 }
                 T{ arc f { "s" "E" } 1 }
                 T{ arc f { "F" "t" } 1 }
                 T{ arc f { "G" "t" } 1 }
                 T{ arc f { "H" "t" } 1 }
                 T{ arc f { "I" "t" } 1 }
                 T{ arc f { "J" "t" } 1 }
                 T{ arc f { "A" "G" } 1 }
                 T{ arc f { "A" "H" } 1 }
                 T{ arc f { "B" "F" } 1 }
                 T{ arc f { "B" "I" } 1 }
                 T{ arc f { "C" "H" } 1 }
                 T{ arc f { "C" "J" } 1 }
                 T{ arc f { "D" "F" } 1 }
                 T{ arc f { "D" "G" } 1 }
                 T{ arc f { "E" "G" } 1 }
                 T{ arc f { "E" "I" } 1 } } } ;

: demo ( -- ) "
  =========================
  | Figure 7.12, Page 372 |
  =========================\n" print
  fig7.12 dup 0 print-network maximize-flow
nl nl "
  =========================
  | Figure 7.21, Page 375 |
  =========================\n" print
  fig7.21 dup 0 print-network maximize-flow
nl nl "
  ============================
  | Exercises 7.2, Problem 2 |
  ============================\n" print
  7.2problem2 dup 0 print-network maximize-flow
nl nl "
  ============================
  | Exercises 7.2, Problem 8 |
  ============================\n" print
  7.2problem8 dup 0 print-network
  increase-flow 1 print-network
nl nl "
  =============================
  | Exercises 7.2, Problem 11 |
  =============================\n" print
  7.2problem11 dup 0 print-network maximize-flow
nl nl "
  ============================
  | Exercises 7.3, Problem 8 |
  ============================\n" print
  7.3problem8 dup 0 print-network
  label-source scan-and-label 1 print-network
nl nl "
  =============================
  | Exercises 7.4, Problem 16 |
  =============================\n" print
  7.4problem16 dup 0 print-network maximize-flow
nl nl "
  =============================
  | Exercises 7.4, Problem 22 |
  =============================\n" print
  7.4problem22 dup 0 print-network maximize-flow ;

MAIN: demo