#|

todo:

  - tree diff

  - step until a particular reduction happens (or a choice point is reached)
  
  - break points: supply a function to traces that is a predicate on terms, indicating if
    this one is one where the -> button should stop.

|#

(module stepper mzscheme
  (require (lib "pretty.ss")
           (lib "mred.ss" "mred")
           (lib "class.ss")
           (lib "framework.ss" "framework")
           (lib "list.ss")
           (lib "graph.ss" "mrlib")
           (lib "contract.ss")
           "size-snip.ss"
           "reduction-semantics.ss")
  
  (provide stepper)
  
  (define dot-spacing 20)
  (define dot-size 10)
  (define initial-color "white")
  (define in-path-color "orchid")
  (define visible-color "purple")
  (define cycle-color "yellow")
  (define visible-cycle-color "gold")
  
  (define (stepper lang red term)
    
    ;; all-nodes-ht : hash-table[sexp -o> (is-a/c node%)]
    (define all-nodes-ht (make-hash-table 'equal))
    
    (define root (new node% 
                      [all-nodes-ht all-nodes-ht]
                      [term term]
                      [red red]
                      [change-path (λ (new-node) (change-path new-node))]))
    
    ;; path : (listof (listof (is-a/c node%))
    ;; the currently visible columns in the pasteboard
    (define path (list (list root)))
    
    (define f (new frame% 
                   [label "PLT Redex Stepper"]
                   [width 600]
                   [height 500]))
    (define dp (new vertical-panel% [parent f]))
    (define hp (new horizontal-panel% [parent dp]))
    (define pb (new columnar-pasteboard% [moved (λ (a b c d) 
                                                  (when (procedure? moved)
                                                    (moved a b c d)))]))
    (define ec (new forward-size-editor-canvas% [parent hp] [editor pb]))
    (define bp-outer (new vertical-panel% [parent hp] [stretchable-width #f]))
    (define bp (new vertical-panel% [parent bp-outer] [stretchable-width #f]))
    (define bp-spacer (new grow-box-spacer-pane% [parent bp-outer]))
    
    (define zoom-out-pb (new zoom-out-pasteboard%))
    (define zoom-out-ec (new editor-canvas% 
                             [stretchable-height #f]
                             [parent dp]
                             [editor zoom-out-pb]))
    
    (define (update-buttons)
      (let ([last-column (car (last-pair path))])
        (let ([last-column last-column])
          (let loop ([children (send bp get-children)]
                     [n 0])
            (cond
              [(= n (length last-column)) 
               (send bp change-children
                     (λ (l)
                       (filter (λ (p) (not (memq p children))) l)))
               (void)]
              [(null? children)
               (new button-object% 
                    [parent bp]
                    [n n])
               (loop children
                     (+ n 1))]
              [else
               (loop (cdr children)
                     (+ n 1))])))
        (let ([button-objects (send bp get-children)])
          (if (null? (cdr button-objects))
              (send (car button-objects) hide-vertical)
              (for-each (λ (x) (send x show-vertical))
                        button-objects))
          (for-each (λ (node button-object)
                      (cond
                        [(not (null? (send node get-cycle)))
                         (send button-object step-goes-back #t)
                         (send button-object enable-step #f)]
                        [else
                         (send button-object step-goes-back #f)
                         (send button-object enable-step (not (null? (send node get-successors))))]))
                    last-column
                    button-objects))))
    
    (define button-object%
      (class vertical-panel%
        (init-field n)
        (super-new [style '(border)]
                   [alignment '(left center)])
        
        (inherit change-children)
        (define/public (hide-vertical)
          (change-children (λ (x) (remq expand-button x))))
        (define/public (show-vertical)
          (change-children (λ (x) (if (memq expand-button x)
                                      x
                                      (append x (list expand-button))))))
        
        (define/public (enable-step on?)
          (send step-button enable on?))
        
        (define/public (step-goes-back back?)
          (send step-button set-label (if back? "↩" "→")))
        
        (define step-button
          (new button%
               [label "→"]
               [callback (λ (x y) (forward-step n))]
               [parent this]))
        (define expand-button
          (new button% 
               [label "↕"] 
               [callback (λ (x y) (expand n))] 
               [parent this]))))
    
    (define (forward-step n)
      (let* ([last-pr (last-pair path)]
             [last-column (car last-pr)]
             [click-target (list-ref last-column n)])
        (cond
          [(not (null? (send click-target get-cycle)))
           (void)]
          [else
           (let ([new-children (begin (send click-target force)
                                      (send click-target get-children))])
             (unless (null? new-children)
               (for-each (λ (x) (send x set-in-path? (eq? x click-target))) last-column)
               (for-each (λ (x) (send x set-in-path? #t)) new-children)
               (set-car! last-pr (list click-target))
               (set-cdr! last-pr (list new-children))
               (send pb begin-edit-sequence)
               (pb-change-columns)
               (pb-last-column-visible)
               (send pb end-edit-sequence)
               (update-buttons)))])))
    
    (define (expand n)
      (let* ([last-pr (last-pair path)]
             [last-column (car last-pr)]
             [survivor (list-ref last-column n)])
        (for-each (λ (x) (send x set-in-path? (eq? x survivor))) last-column)
        (set-car! last-pr (list survivor))
        (send pb begin-edit-sequence)
        (pb-change-columns)
        (pb-last-column-visible)
        (send pb end-edit-sequence)
        (update-buttons)))
    
    (define (moved left top right bottom)
      (let ([bx (box 0)])
        (let loop ([path path])
          (cond
            [(null? path) (void)]
            [else
             (let* ([path-ele (car path)]
                    [snip (send (car path-ele) get-big-snip)]
                    [visible?
                     (or (begin (send pb get-snip-location snip bx #f #f)
                                (<= left (unbox bx) right))
                         (begin (send pb get-snip-location snip bx #f #t)
                                (<= left (unbox bx) right)))])
               (for-each (λ (node) (send node set-visible? visible?))
                         path-ele))
             (loop (cdr path))]))))
    
    (define (get-path-to-root node)
      (let loop ([node node]
                 [acc null])
        (let ([parents (send node get-parents)])
          (cond
            [(null? parents) (cons (list node) acc)]
            [node (loop (car parents) 
                        (cons (list node) acc))]))))
    
    (define (change-path new-node)
      (cond
        [(ormap (λ (l) (memq new-node l)) path)
         ;; if this node is in the current path, just move the view
         (let* ([snip (send new-node get-big-snip)]
                [admin (send snip get-admin)]
                [bw (box 0)]
                [bh (box 0)])
           (send admin get-view #f #f bw bh snip)
           (send pb scroll-to snip 0 0 (max 10 (unbox bw)) (unbox bh) #t 'start))]
        [else
         (let ([new-path (get-path-to-root new-node)])
           (let loop ([new-path new-path]
                      [path path])
             (cond
               [(or (null? path)
                    (null? new-path)
                    (not (equal? (car path) (car new-path))))
                (for-each (λ (old-ele) (for-each (λ (x) (send x set-in-path? #f)) old-ele))
                          path)
                (for-each (λ (old-ele) (for-each (λ (x) (send x set-in-path? #t)) old-ele))
                          new-path)]
               [else
                (loop (cdr new-path) (cdr path))]))
           
           (set! path new-path)
           (send pb begin-edit-sequence)
           (pb-change-columns)
           (pb-last-column-visible)
           (send pb end-edit-sequence)
           (update-buttons))]))
    
    (define (pb-change-columns)
      (send pb change-columns (map (λ (l) (map (λ (x) (send x get-big-snip)) l)) 
                                   path))
      (send zoom-out-pb refresh-tree root))
    
    ;; makes the last column visible
    (define (pb-last-column-visible)
      (for-each
       (λ (x)
         (let* ([s (send x get-big-snip)]
                [w (send pb get-snip-width s)]
                [admin (send s get-admin)])
           (when admin
             (send admin scroll-to s 0 0 w 10 #t))))
       (car (last-pair path))))
  
    (hash-table-put! all-nodes-ht term root)
    (send f show #t)
    (send root set-in-path? #t)
    (pb-change-columns)
    (update-buttons)) 

  (define node%
    (class object%
      (init-field term 
                  red
                  change-path
                  all-nodes-ht)
      (init [parent #f])
      
      (define parents (if parent 
                          (list parent)
                          '()))
      ;; cycle : (listof node)
      ;; the nodes that have the same term as this one, due to a cycle in the reduction graph
      (define cycle '())
      (define children #f)
      (define big-snip (mk-big-snip term this))
      (define dot-snip (new dot-snip% [node this]))
      (define in-path? #f)
      (define visible? #f)
    
      (define successors #f)
      (define/public (get-successors)
        (unless successors
          (set! successors (apply-reduction-relation red term)))
        successors)
      
      (define/public (move-path)
        (change-path this))
      
      (define/public (set-in-path? p?)
        (set! in-path? p?)
        (update-color))
      
      (define/public (set-visible? v?)
        (set! visible? v?)
        (update-color))
      
      (define/private (update-color)
        (send dot-snip set-color 
              (cond
                [(and visible? in-path? (not (null? cycle)))
                 visible-cycle-color]
                [(not (null? cycle))
                 cycle-color]
                [(and visible? in-path?)
                 visible-color]
                [in-path?
                 in-path-color]
                [else
                 initial-color])))
      
      (define/public (get-cycle) cycle)
      (define/public (add-cycle c) (set! cycle (cons c (remq c cycle))))
      (define/public (get-term) term)
      (define/public (get-big-snip) big-snip)
      (define/public (get-dot-snip) dot-snip)
      (define/public (get-parents) parents)
      (define/public (add-parent p)
        (add-links (send p get-dot-snip) dot-snip)
        (set! parents (cons p parents)))
      (define/public (get-children) (or children '()))
      (define/public (force)
        (unless children
          (set! children
                (map (λ (x) (make-child x)) (get-successors)))))
      
      (define/private (make-child term)
        (let ([already-there (hash-table-get all-nodes-ht term #f)]
              [mk-child-node
               (λ ()
                 (new node% 
                      [term term]
                      [red red]
                      [change-path change-path]
                      [all-nodes-ht all-nodes-ht]
                      [parent this]))])
          (cond
            [(and already-there
                  (is-parent? already-there))
             (let ([n (mk-child-node)])
               (send n add-cycle already-there)
               (send already-there add-cycle n)
               n)]
            [already-there
             (send already-there add-parent this)
             already-there]
            [else
             (let ([child-node (mk-child-node)])
               (hash-table-put! all-nodes-ht term child-node)
               child-node)])))
      
        (define/private (is-parent? node)
          (let loop ([parents (get-parents)])
            (ormap (λ (p)
                     (or (eq? p node)
                         (loop (send p get-parents))))
                   parents)))
        
      (super-new)
      (when cycle
        (send dot-snip set-color cycle-color))
      (when parent
        (add-links (send parent get-dot-snip) dot-snip))))
    
  (define zoom-out-pasteboard%
    (class (graph-pasteboard-mixin pasteboard%)
      (inherit insert move-to get-canvas)

      (inherit find-snip set-caret-owner global-to-local)
      (define/override (on-event evt)
        (when (send evt button-down?)
          (let ([x (box (send evt get-x))]
                [y (box (send evt get-y))])
            (global-to-local x y)
            (let ([s (find-snip (unbox x) (unbox y))])
              (when s
                (set-caret-owner s 'immediate)))))
        (super on-event evt))
      
      (define/public (refresh-tree root)
        (let ([level-ht (make-hash-table)]
              [node-to-level-ht (make-hash-table)]
              [max-n 0])
          
          (let loop ([tree root]
                     [n 0])
            (let ([old-level (hash-table-get node-to-level-ht tree #f)])
              (cond
                [(not old-level)
                 (hash-table-put! node-to-level-ht tree n)
                 (hash-table-put! level-ht n (cons tree (hash-table-get level-ht n '())))]
                [(< old-level n)
                 (hash-table-put! level-ht old-level (remq tree (hash-table-get level-ht old-level)))
                 (hash-table-put! level-ht n (cons tree (hash-table-get level-ht n '())))
                 (hash-table-put! node-to-level-ht tree n)]
                [else
                 (void)])
              (set! max-n (max n max-n))
              (for-each (λ (x) (loop x (+ n 1))) (send tree get-children))))
          
          (let ([tallest-column (apply max (hash-table-map level-ht (λ (x y) (length y))))])
            
            (let loop ([n 0])
              (when (<= n max-n)
                (let ([nodes (reverse (hash-table-get level-ht n))])
                  (let loop ([nodes nodes]
                             [y (/ (- (* tallest-column dot-spacing)
                                      (* (length nodes) dot-spacing)) 
                                   2)])
                    (cond
                      [(null? nodes) (void)]
                      [else
                       (let* ([node (car nodes)]
                              [dot-snip (send node get-dot-snip)])
                         (insert dot-snip (* n dot-spacing) y)   ;; in case the snip's been inserted already
                         (move-to dot-snip (* n dot-spacing) y)  ;; also do the move to
                         (loop (cdr nodes) (+ y dot-spacing)))])))
                (loop (+ n 1))))
            
            (let ([canvas (get-canvas)])
              (send canvas min-client-height (* tallest-column dot-spacing))))))
      (super-new)
      (inherit set-draw-arrow-heads?)
      (set-draw-arrow-heads? #f)))

  (define (set-box/f b v) (when (box? b) (set-box! b v)))
  
  (define dot-snip%
    (class (graph-snip-mixin snip%)
      (init-field node)
      
      (inherit get-admin)
      (define color initial-color)
      (define/public (set-color c)
        (unless (equal? color c)
          (set! color c)
          (let ([admin (get-admin)])
            (when admin
              (send admin needs-update this 0 0 dot-size dot-size)))))
      (define/override (get-extent dc x y wb hb descentb spaceb lspaceb rspaceb)
        (set-box/f wb dot-size)
        (set-box/f hb dot-size)
        (set-box/f descentb 0)
        (set-box/f spaceb 0)
        (set-box/f lspaceb 0)
        (set-box/f rspaceb 0))
      (define/override (draw dc x y left top right bottom dx dy draw-caret)
        (let ([smoothing (send dc get-smoothing)]
              [brush (send dc get-brush)])
          (send dc set-smoothing 'aligned)
          (send dc set-brush color 'solid)
          (send dc draw-ellipse x y dot-size dot-size)
          (send dc set-brush brush)
          (send dc set-smoothing smoothing)))
      (define/override (on-event dc x y editorx editory evt)
        (when (send evt button-up?)
          (send node move-path)))
      
      (define/override (copy) (new snip%))
      (super-new)
      (inherit set-snipclass set-flags get-flags)
      (set-flags (cons 'handles-events (get-flags)))
      (set-snipclass dot-snipclass)))
  
  (define dot-snipclass
    (new
     (class snip-class%
       (define/override (read f)
         (new dot-snip%))
       (super-new))))
  (send dot-snipclass set-classname "plt-redex:dot")
  (send dot-snipclass set-version 1)
  (send (get-the-snip-class-list) add dot-snipclass)
       
  
  (define forward-size-editor-canvas%
    (class editor-canvas%
      (inherit get-editor)
      (define/override (on-size w h)
        (send (get-editor) update-heights))
      (super-new)))
  
  (define (mk-big-snip sexp node)
    (let* ([txt (new scheme:text%)]
           [s (new big-snip% 
                   [node node]
                   [editor txt]
                   [expr sexp]
                   [char-width 40]
                   [pp default-pretty-printer])])
      (send txt set-autowrap-bitmap #f)
      (send txt freeze-colorer)
      (send s format-expr)
      s))
  
  (define big-snip%
    (class size-editor-snip%
      (init-field node)
      (define/public (get-node) node)
      (super-new)))
  
  (define columnar-pasteboard% 
    (class (resizing-pasteboard-mixin pasteboard%)
      (init-field moved)
      
      (define current-columns '())
      (inherit insert remove find-snip)
      
      ;; strange to think that this is the way to catch
      ;; different snips becoming visible in the editor, but oh well.
      (define/override (on-paint before? dc left top right bottom dx dy draw-caret)
        (super on-paint before? dc left top right bottom dx dy draw-caret)
        (unless before?
          (let ([admin (get-admin)])
            (when admin
              (let ([bx (box 0)]
                    [by (box 0)]
                    [bw (box 0)]
                    [bh (box 0)])
                (send admin get-view bx by bw bh)
                (moved (unbox bx)
                       (unbox by)
                       (+ (unbox bx) (unbox bw))
                       (+ (unbox by) (unbox bh))))))))
      
      (define/public (change-columns orig-new-columns)
        (let loop ([current-columns current-columns]
                   [new-columns orig-new-columns])
          (cond
            [(and (null? current-columns)
                  (null? new-columns))
             (void)]
            [(null? new-columns) 
             (insert/remove current-columns '())]
            [(null? current-columns)
             (insert/remove '() new-columns)]
            [(equal? (car current-columns)
                     (car new-columns))
             (loop (cdr current-columns)
                   (cdr new-columns))]
            [else
             (insert/remove current-columns new-columns)]))
        (set! current-columns orig-new-columns)
        (update-heights))
      
      ;; insert/remove : (listof (listof snip)) (listof (listof snip)) -> void
      (define/private (insert/remove to-remove to-insert)
        (let ([flat-to-remove (apply append to-remove)]
              [flat-to-insert (apply append to-insert)])
          (for-each
           (λ (x) (unless (memq x flat-to-insert)
                    (remove x)))
           flat-to-remove)
          (for-each (λ (x) (insert x)) flat-to-insert)))
      
      (inherit get-admin move-to resize)
      (define/public (update-heights)
        (let ([admin (get-admin)])
          (let-values ([(w h) (get-view-size)])
            (let loop ([columns current-columns]
                       [x 0])
              (cond
                [(null? columns) (void)]
                [else
                 (let* ([column (car columns)]
                        [base-space (quotient h (length column))]
                        [widest
                         (let loop ([snips column]
                                    [extra-space (modulo h (length column))]
                                    [y 0]
                                    [widest 0])
                           (cond
                             [(null? snips) widest]
                             [else 
                              (let* ([snip (car snips)]
                                     [sw (get-snip-width snip)]
                                     [h (+ base-space
                                           (if (zero? extra-space)
                                               0
                                               1))])
                                (move-to snip x y)
                                (resize snip sw h)
                                (loop (cdr snips)
                                      (if (zero? extra-space)
                                          0
                                          (- extra-space 1))
                                      (+ y h)
                                      (max widest sw)))]))])
                   (loop (cdr columns)
                         (+ x widest)))])))))

      (inherit get-snip-location)
      (define/public (get-snip-width snip)
        (let ([lb (box 0)]
              [rb (box 0)])
          (get-snip-location snip lb #f #f)
          (get-snip-location snip rb #f #t)
          (- (unbox rb) (unbox lb))))
      
      (define/private (get-view-size)
        (let ([admin (get-admin)])
          (if admin
              (let ([wb (box 0)]
                    [hb (box 0)])
                (send admin get-view #f #f wb hb)
                (values (unbox wb) (- (unbox hb) 2)))
              (values 10 10))))
      
      (super-new)))
  
  (define (all-but-last l)
    (let loop ([l l])
      (cond
        [(null? (cdr l)) null]
        [else (cons (car l) (loop (cdr l)))])))
  
  (define (record-differences sexp1 sexp2)
    (let ([ht (make-hash-table)])
      (let loop ([sexp1 sexp1]
                 [sexp2 sexp2])
        (cond
          [(eq? sexp1 sexp2) (void)]
          [(and (pair? sexp1)
                (pair? sexp2)
                (equal? (d-length sexp1)
                        (d-length sexp2)))
           (for-each/d loop sexp1 sexp2)]
          [(equal? sexp1 sexp2) (void)]
          [else 
           (hash-table-put! ht sexp1 #t)
           (hash-table-put! ht sexp2 #t)]))
      ht))
  
  (define (show-differences s1 s2)
    (define diff-ht (record-differences s1 s2))
    (define f (new frame% [label ""] [width 600] [height 500]))
    (define hp (new horizontal-panel% [parent f]))
    (define t1 (new text:basic%))
    (define t2 (new text:basic%))
    (define c1 (new editor-canvas% 
                    [parent hp]
                    [editor t1]))
    (define c2 (new editor-canvas% 
                    [parent hp]
                    [editor t2]))
    (render-sexp/colors s1 diff-ht t1)
    (render-sexp/colors s2 diff-ht t2)
    (send f show #t))
  
  ;; render-sexp/colors : sexp ht text -> void
  (define (render-sexp/colors sexp diff-ht text)
    (let ([start #f])
      (parameterize ([pretty-print-columns 30]
                     [pretty-print-pre-print-hook
                      (λ (obj port)
                        (when (hash-table-get diff-ht obj #f)
                          (set! start (send text last-position))))]
                     [pretty-print-post-print-hook
                      (λ (obj port)
                        (when (hash-table-get diff-ht obj #f)
                          (send text highlight-range
                                start
                                (send text last-position)
                                (send the-color-database find-color "magenta"))
                          (set! start #f)))])
        (pretty-print sexp (open-output-text-editor text)))
      (send text change-style 
            (make-object style-delta% 'change-family 'modern)
            0
            (send text last-position))))
  
  (define (for-each/d f l1 l2)
    (let loop ([l1 l1]
               [l2 l2])
      (cond
        [(pair? l1)
         (f (car l1) (car l2))
         (loop (cdr l1) (cdr l2))]
        [(null? l1) (void)]
        [else (f l1 l2)])))
  
  (define (d-length l1)
    (let loop ([l1 l1]
               [n 0])
      (cond
        [(pair? l1) (loop (cdr l1) (+ n 1))]
        [(null? l1) n]
        [else (cons 'dotted (+ n 1))]))))