#lang racket

;;; dds/generic

;;; Provides the definition of several generic interfaces for discrete
;;; dynamical systems.

(require racket/generic graph "utils.rkt")

(provide
 ;; Generics
 gen:dds
 ;; Functions
 (contract-out [dds-step-one (-> dds? any/c (set/c any/c))]
               [dds-step-one-annotated (-> dds? any/c (set/c (cons/c any/c any/c)))]
               [dds-step (-> dds? (set/c any/c #:kind 'dont-care) (set/c any/c))]
               [dds-build-state-graph (-> dds? (set/c any/c #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph (-> dds? (set/c any/c #:kind 'dont-care) number? graph?)]
               [dds-build-state-graph-annotated (-> dds? (set/c any/c #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph-annotated (-> dds? (set/c any/c #:kind 'dont-care) number? graph?)])
 ;; Predicates
 (contract-out [dds? (-> any/c boolean?)]))

;;; Given a dds and a set of starting states, produce the set of
;;; states reachable in one step.  This is a fallback for dds-step.
(define (fallback-dds-step dds ss)
  (apply set-union (for/list ([s ss]) (dds-step-one dds s))))

;;; Given a dds, a set of starting states, and a range whose length
;;; determines the number of steps to run, produces the edges (and the
;;; edge labels) of the state graph reachable from the starting states
;;; in this many steps.  The last argument is a function similar do
;;; dds-step-one-annotated: given a state, it should produce the set
;;; of next states, labelled with appropriate labels.
;;;
;;; This is a fallback for dds-build-state-graph,
;;; dds-build-n-step-state-graph, dds-build-state-graph-annotated, and
;;; dds-build-n-step-state-graph-annotated.
(define (fallback-dds-build-state-graph-edges dds states step-range step-func)
  (for/fold ([edges empty]
             [labels empty]
             [current-states states]
             [visited-states states]
             #:result (collect-by-key/sets edges labels))
            ([i step-range]
             #:break (set-empty? current-states))
    (for/fold ([new-edges empty]
               [new-labels empty]
               [new-states (set)]
               #:result (values
                         (append edges new-edges)
                         (append labels new-labels)
                         (set-subtract new-states visited-states)
                         (set-union current-states visited-states)))
              ([s current-states])
      (for/fold ([edges-to-add empty]
                 [labels-to-add empty]
                 [states-to-add empty]
                 #:result (values
                           (append new-edges edges-to-add)
                           (append new-labels labels-to-add)
                           (set-union (list->set states-to-add) new-states)))
                ([next (set->list (step-func dds s))])
        (match next
          [(cons label s-next)
           (values
            (cons (list s s-next) edges-to-add)
            (cons label labels-to-add)
            (cons s-next states-to-add))])))))

;;; Run dds-step-one, and produce a set of new states annotated with a
;;; 'dummy annotation.
(define (dummy-annotated-dds-step-one dds s)
  (for/set ([new-s (dds-step-one dds s)])
    (cons 'dummy new-s)))

;;; A poor man's fallback for dds-step-one-annotated.
(define (fallback-dds-step-one-annotated dds s) (dds-step-one-annotated dds s))

;;; Run dds-step-one-annotated, but then discard the annotations.
;;; This is a fallback for dds-step-one.
(define (fallback-dds-step-one dds state)
  (for/set ([annotation-state (dds-step-one-annotated dds state)])
    (match annotation-state
      [(cons a s) s])))

;;; A discrete dynamical system.
(define-generics dds
  ;; Given a dds and a state, produce the next states of the dds.
  ;; This method falls back to calling dds-step-one-annotated,
  ;; discarding the annotations.
  (dds-step-one dds state)
  ;; Given a dds and a state, produce the next states paired with some
  ;; annotation.  Typical usage would include including the
  ;; information about the update mode.  This method has no fallback.
  (dds-step-one-annotated dds state)
  ;; Given a dds and a set of starting states, produce the set of
  ;; states reachable in one step.  This method falls back to running
  ;; dds-step-one for all states.
  (dds-step dds states)
  ;; Given a dds and a set of starting states, produces the state
  ;; graph reachable from the starting states.  This method falls back
  ;; to exploring the state graph with dds-step-one.
  (dds-build-state-graph dds states)
  ;; Given a dds, a set of starting states, and a number of steps to
  ;; run, produces the state graph reachable from the starting states
  ;; in this many steps.  This method falls back to exploring the
  ;; state graph with dds-step-one.
  (dds-build-n-step-state-graph dds states nsteps)
  ;; Given a dds and a set of starting states, produces the labelled
  ;; state graph reachable from the starting states.  This method
  ;; falls back to exploring the state graph with
  ;; dds-step-one-annotated.
  (dds-build-state-graph-annotated dds states)
  ;; Given a dds, a set of starting states, and a number of steps to
  ;; run, produces the labelled state graph reachable from the
  ;; starting states in this many steps.  This method falls back to
  ;; exploring the state graph with dds-step-one-annotated.
  (dds-build-n-step-state-graph-annotated dds states nsteps)

  #:defined-predicate dds-implements?
  #:fallbacks
  [(define dds-step fallback-dds-step)
   (define dds-step-one fallback-dds-step-one)

   ;; Run fallback-dds-build-state-graph-edges with an infinite range,
   ;; which will make it stop only when it has explored all the state
   ;; graph.  Use dummy-annotated-dds-step-one to produce dummy edge
   ;; labels and then discard them.
   (define (dds-build-state-graph dds states)
     (let-values ([(edges labels)
                   (fallback-dds-build-state-graph-edges
                    dds states (in-naturals)
                    dummy-annotated-dds-step-one)])
       (directed-graph edges)))

   ;; Run fallback-dds-build-state-graph-edges within the given range.
   ;; Use dummy-annotated-dds-step-one to produce dummy edge labels
   ;; and then discard them.
   (define (dds-build-n-step-state-graph dds states nsteps)
     (let-values ([(edges labels)
                   (fallback-dds-build-state-graph-edges
                    dds states (in-range nsteps)
                    dummy-annotated-dds-step-one)])
       (directed-graph edges)))

   (define (dds-build-state-graph-annotated dds states)
     (let-values ([(edges labels)
                   (fallback-dds-build-state-graph-edges
                    dds states (in-naturals)
                    fallback-dds-step-one-annotated)])
       (directed-graph edges labels)))

   (define (dds-build-n-step-state-graph-annotated dds states nsteps)
     (let-values ([(edges labels)
                   (fallback-dds-build-state-graph-edges
                    dds states (in-range nsteps)
                    fallback-dds-step-one-annotated)])
       (directed-graph edges labels)))])