#lang typed/racket

(require "utils.rkt" typed/graph)

(provide dds%)

(define dds%
  (class object%
    #:forall (State Modality)
    (super-new)

    (: step (-> State (Listof State)))
    (define/public (step st)
      (map (inst cdr Modality State) (step/annotated st)))

    (: step/annotated (-> State (Listof (Pairof Modality State))))
    (define/abstract/error (step/annotated st))

    (: step* (-> (Listof State) (Listof State)))
    (define/public (step* sts)
      (remove-duplicates
       (apply append
              (for/list : (Listof (Listof State)) ([s sts])
                (step s)))))

    (: build-state-graph (-> (Listof State) Graph))
    (define/public (build-state-graph sts)
      (build-state-graph* sts 'full))

    (: build-state-graph/annotated (-> (Listof State) Graph))
    (define/public (build-state-graph/annotated sts)
      (build-state-graph*/annotated sts 'full))

    (: build-state-graph* (-> (Listof State) (U Positive-Integer 'full) Graph))
    (define/public (build-state-graph* sts nsteps)
      (unweighted-graph/directed
       (assert-type (get-edges (build-state-graph*/annotated sts nsteps))
                    (Listof (List Any Any)))))

    (: build-state-graph*/annotated (-> (Listof State) (U Positive-Integer 'full) Graph))
    (define/public (build-state-graph*/annotated sts nsteps)
      (define (all-steps-done? k)
        (if (equal? nsteps 'full)
            #f ; keep going forever
            (>= (assert-type k Integer) nsteps)))
      (weighted-graph/directed
       (let build-graph : (Listof (List Modality State State))
            ([visited-states : (Setof State) (set)]
             [states : (Setof State) (list->set sts)]
             [edges : (Listof (List Modality State State)) '()]
             [k 1])
         (define new-edges
           (for*/list : (Listof (List Modality State State))
                      ([s (in-set states)] ; the state we are looking at
                       [out (in-list (step/annotated s))]) ; the arrows going out of s
             (list (car out) s (cdr out))))
         (define edges/full (append edges new-edges))
         (define new-states
           (list->set (map (inst caddr Modality State State) new-edges)))
         (define new-states-pruned
           (set-subtract new-states visited-states))
         (if (or (set-empty? new-states-pruned) (all-steps-done? k))
             edges/full
             (build-graph (set-union visited-states new-states-pruned)
                          new-states-pruned
                          edges/full
                          (add1 k))))))
    ))