(module suffixtree mzscheme
  (require (prefix u: "private/ukkonen2.ss")
           (prefix l: "private/label.ss")
           (prefix s: "private/structs.ss"))
  (require (lib "contract.ss"))


  ;; Basic predicates
  (provide/contract (tree? (-> any/c boolean?)))
  (define tree? s:suffix-tree?)

  (provide/contract (node? (-> any/c boolean?)))
  (define node? s:node?)

  (provide/contract (label? (-> any/c boolean?)))
  (define label? l:label?)

  (provide/contract (label-element? (-> any/c boolean?)))
  (define label-element? l:label-element?)


  ;; Labels
  (define nonempty-label/c
    (flat-named-contract "nonempty-label"
                         (lambda (datum)
                           (and (l:label? datum)
                                (not (l:label-empty? datum))))))

  (provide/contract (string->label (-> string? label?)))
  (define string->label l:string->label)

  (provide/contract (string->label/with-sentinel (-> string? label?)))
  (define string->label/with-sentinel l:string->label/with-sentinel)

  (provide/contract (vector->label (-> vector? label?)))
  (define vector->label l:vector->label)

  (provide/contract (vector->label/with-sentinel (-> vector? label?)))
  (define vector->label/with-sentinel l:vector->label/with-sentinel)

  (provide/contract (label->vector (-> label? vector?)))
  (define label->vector l:label->vector)
  
  (provide/contract (label->string (-> label? string?)))
  (define label->string l:label->string)
                     
  (provide/contract (label-equal? (-> label? label? boolean?)))
  (define label-equal? l:label-equal?)

  (provide/contract (label-length (-> label? natural-number/c)))
  (define label-length l:label-length)

  (provide/contract (label-ref (-> label? natural-number/c label-element?)))
  (define label-ref l:label-ref)

  (provide/contract (sublabel (case->
                               (-> label? natural-number/c natural-number/c label?)
                               (-> label? natural-number/c label?))))
  (define sublabel l:sublabel)

  (provide/contract (label-source-id (-> label? integer?)))
  (define label-source-id l:label-source-id)

  (provide/contract (label-source-eq? (-> label? label? boolean?)))
  (define label-source-eq? l:label-same-source?)
  

  
  ;; Trees
  (provide/contract (make-tree (-> tree?)))
  (define make-tree s:new-suffix-tree)
  
  (provide/contract (tree-root (-> tree? node?)))
  (define tree-root s:suffix-tree-root)

  (provide/contract (tree-add! (-> tree? nonempty-label/c void?)))
  (define tree-add! u:suffix-tree-add!)
  


  ;; Nodes
  (provide/contract (node-up-label (-> node? label?)))
  (define node-up-label s:node-up-label)

  (provide/contract (node-children (-> node? (listof node?))))
  (define node-children s:node-children)

  (provide/contract (node-parent (-> node? (union false/c node?))))
  (define node-parent s:node-parent)

  (provide/contract (node-find-child (-> node? label-element?
                                         (union false/c node?))))
  (define node-find-child s:node-find-child)

  (provide/contract (node-suffix-link (-> node? (union false/c node?))))
  (define node-suffix-link s:node-suffix-link)

  
  ;; tree-walk: tree label (node number -> A)
  ;;                       (node number number -> B) -> (union A B)
  (provide/contract (tree-walk (-> tree? label? procedure? procedure?
                                   any/c)))
  (define (tree-walk tree input-label succeed-f fail-f)
    (letrec [(comparing-label-elements
              (lambda (node up-label up-label-offset input-label-offset)
                (cond
                 ((= input-label-offset (label-length input-label))
                  (succeed-f node up-label-offset))
                 ((= up-label-offset (label-length up-label))
                  (choosing-next-edge node input-label-offset))
                 ((l:label-element-equal?
                   (label-ref up-label up-label-offset)
                   (label-ref input-label input-label-offset))
                  (comparing-label-elements node up-label
                                            (add1 up-label-offset)
                                            (add1 input-label-offset)))
                 (else
                  (fail-f node up-label-offset input-label-offset)))))
             (choosing-next-edge
              (lambda (node input-label-offset)
                (let ((child
                       (node-find-child
                        node (label-ref input-label input-label-offset))))
                  (if child
                      (comparing-label-elements child
                                                (node-up-label child) 0
                                                input-label-offset)
                      (fail-f node (label-length (node-up-label node))
                              input-label-offset)))))]
      (if (= 0 (label-length input-label))
          (begin
            (succeed-f (tree-root tree) 0))
          ;; Start off the walk at the root.
          (choosing-next-edge (tree-root tree) 0))))


  (provide/contract (tree-contains? (-> tree? label? boolean?)))
  (define (tree-contains? tree label)
    (tree-walk tree label 
               (lambda (node edge-label-offset) #t)
               (lambda (node edge-label-offset input-label-offset) #f)))
  
    
                                    

  
  )