#lang racket

(module typed typed/racket
  (require typed/graph "utils.rkt" "dynamics.rkt")

  (provide
   Species (struct-out reaction) Reaction ReactionName ReactionSystem
   make-reaction enabled? list-enabled union-products apply-rs

   str-triple->reaction ht-str-triples->rs read-org-rs read-context-sequence
   reaction->str-triple rs->ht-str-triples

   (struct-out state) State
   )

  (module+ test
    (require typed/rackunit))

  (define-type Species Symbol)

  (struct reaction ([reactants : (Setof Species)]
                    [inhibitors : (Setof Species)]
                    [products : (Setof Species)])
    #:transparent
    #:type-name Reaction)

  (define-type ReactionName Symbol)

  (: make-reaction (-> (Listof Species) (Listof Species) (Listof Species) Reaction))
  (define (make-reaction r i p) (reaction (list->set r)
                                          (list->set i)
                                          (list->set p)))
  (module+ test
    (test-case "make-reaction"
      (check-equal? (make-reaction '(a b) '(c d) '(e f))
                    (reaction (set 'b 'a) (set 'c 'd) (set 'f 'e)))))

  (: enabled? (-> Reaction (Setof Species) Boolean))
  (define/match (enabled? r s)
    [((reaction r i _) s)
     (and (subset? r s) (set-empty? (set-intersect i s)))])

  (module+ test
    (test-case "enabled?"
      (check-true (enabled? (make-reaction '(a b) '(c d) '())
                            (set 'a 'b 'e)))
      (check-false (enabled? (make-reaction '(a b) '(c d) '())
                             (set 'a 'b 'c)))
      (check-false (enabled? (make-reaction '(a b) '(c d) '())
                             (set 'b 'e)))))

  (define-type ReactionSystem (HashTable ReactionName Reaction))

  (: list-enabled (-> ReactionSystem (Setof Species) (Listof ReactionName)))
  (define (list-enabled rs s)
    (for/list ([(name reaction) (in-hash rs)]
               #:when (enabled? reaction s))
      name))

  (module+ test
    (test-case "list-enabled"
      (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                       'b (make-reaction '(x y) '() '(z))))
      (check-equal? (list-enabled rs (set 'x 'y)) '(b))
      (check-equal? (list-enabled rs (set 'x)) '(a))))

  (: union-products (-> ReactionSystem (Listof ReactionName) (Setof Species)))
  (define (union-products rs as)
    (cond
      [(empty? as) (set)]
      [else (define products (for/list : (Listof (Setof Species))
                                       ([a as])
                               (reaction-products (hash-ref rs a))))
            (apply set-union (assert-type products (NonemptyListof (Setof Species))))]))

  (module+ test
    (test-case "union-products"
      (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                       'b (make-reaction '(x y) '() '(t))))
      (check-equal? (union-products rs '(a b))
                    (set 't 'z))
      (check-equal? (union-products rs '(a))
                    (set 'z))
      (check-equal? (union-products rs '())
                    (set))))

  (: apply-rs (-> ReactionSystem (Setof Species) (Setof Species)))
  (define (apply-rs rs s)
    (let ([as (list-enabled rs s)])
      (union-products rs as)))

  (module+ test
    (test-case "apply-rs"
      (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                       'b (make-reaction '(x y) '() '(t))))
      (check-equal? (apply-rs rs (set 'x 'y))
                    (set 't))
      (check-equal? (apply-rs rs (set 'x))
                    (set 'z))))

  (: str-triple->reaction (-> (List String String String) Reaction))
  (define/match (str-triple->reaction lst)
    [((list str-reactants str-inhibitors str-products))
     (reaction (list->set (read-symbol-list str-reactants))
               (list->set (read-symbol-list str-inhibitors))
               (list->set (read-symbol-list str-products)))])

  (module+ test
    (test-case "str-triple->reaction"
      (check-equal? (str-triple->reaction '("a b" "c d" "e f"))
                    (reaction (set 'b 'a) (set 'c 'd) (set 'f 'e)))))

  (: ht-str-triples->rs (-> (HashTable ReactionName (List String String String))
                            ReactionSystem))
  (define (ht-str-triples->rs ht)
    (for/hash : (HashTable ReactionName Reaction)
              ([(a triple) (in-hash ht)])
      (values a (str-triple->reaction triple))))

  (module+ test
    (test-case "ht-str-triples->rs"
      (check-equal? (ht-str-triples->rs (hash 'a (list "x y" "" "k i")
                                              'b (list "" "x y" "t j")))
                    (hash 'a (reaction (set 'y 'x) (set) (set 'k 'i))
                          'b (reaction (set) (set 'y 'x) (set 't 'j))))))

  (: read-org-rs (-> String ReactionSystem))
  (define (read-org-rs str)
    (ht-str-triples->rs
     (assert-type (read-org-variable-mapping str)
                  (Immutable-HashTable ReactionName (List String String String)))))

  (module+ test
    (test-case "read-org-rs"
      (check-equal?
       (read-org-rs "((\"a\" \"x t\" \"y\" \"z\") (\"b\" \"x\" \"q\" \"z\"))")
       (hash 'a (reaction (set 't 'x) (set 'y) (set 'z))
             'b (reaction (set 'x) (set 'q) (set 'z))))))

  (: read-context-sequence (-> String (Listof (Setof Species))))
  (define (read-context-sequence str)
    (for/list ([sexp (in-list (flatten (string->any str)))])
      (list->set (read-symbol-list (assert-type sexp String)))))

  (module+ test
    (test-case "read-context-sequence"
      (check-equal? (read-context-sequence "((\"x y\") (\"z\") (\"\") (\"t\"))")
                    (list (set 'x 'y) (set 'z) (set) (set 't)))))

  (: reaction->str-triple (-> Reaction (Listof String)))
  (define/match (reaction->str-triple r)
    [((reaction r i p))
     (for/list ([c (in-list (list r i p))])
       (drop-first-last (any->string (set->list c))))])

  (module+ test
    (test-case "reaction->str-triple"
      (check-equal? (reaction->str-triple (make-reaction '(x y) '(z t) '(k i)))
                    '("x y" "z t" "i k"))))

  (: rs->ht-str-triples (-> ReactionSystem (HashTable ReactionName (Listof String))))
  (define (rs->ht-str-triples rs)
    (for/hash : (HashTable ReactionName (Listof String))
              ([(a r) (in-hash rs)])
      (values a (reaction->str-triple r))))

  (module+ test
    (test-case "rs->ht-str-triples"
      (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                       'b (make-reaction '(x y) '() '(t))))
      (rs->ht-str-triples rs)))

  (struct state ([result : (Setof Species)]
                 [rest-contexts : (Listof (Setof Species))])
    #:transparent
    #:type-name State)
  )

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

(provide
 ;; Structures
 (struct-out reaction)
 (struct-out state)
 (struct-out dynamics)
 ;; Functions
 (contract-out [enabled? (-> reaction? (set/c symbol?) boolean?)]
               [list-enabled (-> reaction-system/c (set/c species?) (listof symbol?))]
               [union-products (-> reaction-system/c (listof symbol?) (set/c species?))]
               [apply-rs (-> reaction-system/c (set/c species?) (set/c species?))]
               [ht-str-triples->rs (-> (hash/c symbol? (list/c string? string? string?)) reaction-system/c)]
               [read-org-rs (-> string? reaction-system/c)]
               [read-context-sequence (-> string? (listof (set/c species?)))]
               [rs->ht-str-triples (-> reaction-system/c (hash/c symbol? (list/c string? string? string?)))]
               [dds-step-one (-> dynamics? state? (set/c state?))]
               [dds-step-one-annotated (-> dynamics? state? (set/c (cons/c (set/c symbol?) state?)))]
               [dds-step (-> dynamics? (set/c state? #:kind 'dont-care) (set/c state?))]
               [dds-build-state-graph (-> dynamics? (set/c state? #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph (-> dynamics? (set/c state? #:kind 'dont-care) number? graph?)]
               [dds-build-state-graph-annotated (-> dynamics? (set/c state? #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph-annotated (-> dynamics? (set/c state? #:kind 'dont-care) number? graph?)]
               [build-interactive-process-graph (-> reaction-system/c (listof (set/c species?)) graph?)]
               [build-reduced-state-graph (-> reaction-system/c (listof (set/c species?)) graph?)]
               [pretty-print-reduced-state-graph (-> graph? graph?)]
               [build-interactive-process (-> reaction-system/c (listof (set/c species?)) (listof (list/c (set/c species?) (set/c species?))))]
               [pretty-print-state-graph (-> graph? graph?)])
 ;; Predicates
 (contract-out [species? (-> any/c boolean?)])
 ;; Contracts
 (contract-out [reaction-system/c contract?]))

(module+ test
  (require rackunit))


;;; =================
;;; Basic definitions
;;; =================

;;; A species is a symbol.
(define species? symbol?)

;;; A reaction is a triple of sets, giving the reactants, the
;;; inhibitors, and the products, respectively.
(struct reaction (reactants inhibitors products) #:transparent)

;;; A reaction is enabled on a set if all of its reactants are in the
;;; set and none of its inhibitors are.
(define/match (enabled? r s)
  [((reaction r i p) s)
   (and (subset? r s) (set-empty? (set-intersect i s)))])

;;; A reaction system is a dictionary mapping reaction names to
;;; reactions.
(define reaction-system/c (hash/c symbol? reaction?))

;;; Returns the list of reaction names enabled on a given set.
(define (list-enabled rs s)
  (for/list ([(name reaction) (in-hash rs)]
             #:when (enabled? reaction s))
    name))

;;; Returns the union of the product sets of the given reactions in a
;;; reaction system.  If no reactions are supplied, returns the empty
;;; set.
;;;
;;; This function can be seen as producing the result of the
;;; application of the given reactions to a set.  Clearly, it does not
;;; check whether the reactions are actually enabled.
(define (union-products rs as)
  (if (empty? as)
      (set)
      (apply set-union
             (for/list ([a as])
               (reaction-products (hash-ref rs a))))))

;;; Applies a reaction system to a set.
(define (apply-rs rs s)
  (let ([as (list-enabled rs s)])
    (union-products rs as)))

(module+ test
  (test-case "Basic definitions"
    (define r1 (reaction (set 'x) (set 'y) (set 'z)))
    (define r2 (reaction (set 'x) (set) (set 'y)))
    (define rs (make-immutable-hash (list (cons 'a r1) (cons 'b r2))))
    (define s1 (set 'x 'z))
    (define s2 (set 'x 'y))
    (check-true (enabled? r1 s1))
    (check-false (enabled? r1 s2))
    (check-equal? (list-enabled rs s1) '(a b))
    (check-equal? (list-enabled rs s2) '(b))
    (check-equal? (union-products rs '(a b)) (set 'y 'z))
    (check-equal? (apply-rs rs s1) (set 'y 'z))
    (check-equal? (apply-rs rs s2) (set 'y))))


;;; ====================
;;; Org-mode interaction
;;; ====================

;;; This section contains some useful primitives for Org-mode
;;; interoperability.

;;; Converts a triple of strings to a reaction.
(define/match (str-triple->reaction lst)
  [((list str-reactants str-inhibitors str-products))
   (reaction (list->set (read-symbol-list str-reactants))
             (list->set (read-symbol-list str-inhibitors))
             (list->set (read-symbol-list str-products)))])

;;; Converts a hash table mapping reaction names to triples of strings
;;; to a reaction system.
(define (ht-str-triples->rs ht)
  (for/hash ([(a triple) (in-hash ht)])
    (values a (str-triple->reaction triple))))

(module+ test
  (test-case "ht-str-triples->rs"
    (check-equal?
     (ht-str-triples->rs #hash((a . ("x t" "y" "z"))))
     (make-immutable-hash (list (cons 'a (reaction (set 'x 't) (set 'y) (set 'z))))))))

;;; Reads a reaction system from an Org-mode style string.
(define read-org-rs (compose ht-str-triples->rs read-org-variable-mapping))

(module+ test
  (test-case "read-org-rs"
    (check-equal? (read-org-rs "((\"a\" \"x t\" \"y\" \"z\") (\"b\" \"x\" \"q\" \"z\"))")
                  (hash
                   'a
                   (reaction (set 'x 't) (set 'y) (set 'z))
                   'b
                   (reaction (set 'x) (set 'q) (set 'z))))))

;;; Reads a context sequence from an Org sexp corresponding to a list.
(define (read-context-sequence str)
  (map (compose list->set read-symbol-list) (flatten (string->any str))))

(module+ test
  (test-case "read-context-sequence"
    (check-equal? (read-context-sequence "((\"x y\") (\"z\") (\"\") (\"t\"))")
                  (list (set 'x 'y) (set 'z) (set) (set 't)))))

;;; Converts a reaction to a triple of strings.
(define/match (reaction->str-triple r)
  [((reaction r i p))
   (map (compose drop-first-last any->string set->list)
        (list r i p))])

;;; Converts a reaction system to a hash table mapping reaction names
;;; to triples of strings.
(define (rs->ht-str-triples rs)
  (for/hash ([(a r) (in-hash rs)])
    (values a (reaction->str-triple r))))

(module+ test
  (test-case "rs->ht-str-triples"
    (check-equal?
     (rs->ht-str-triples (make-immutable-hash (list (cons 'a (reaction (set 'x 't) (set 'y) (set 'z))))))
     #hash((a . ("t x" "y" "z"))))))


;;; ============================
;;; Dynamics of reaction systems
;;; ============================

;;; An interactive process of a reaction system is a sequence of
;;; states driven by a sequence of contexts in the following way.
;;; The reaction system starts with the initial context.  Then, at
;;; every step, the result of applying the reaction system is merged
;;; with the next element of the context sequence, and the reaction
;;; system is then applied to the result of the union.  If the
;;; sequence of contexts is empty, the reaction system cannot evolve.

;;; A state of a reaction system is a set of species representing the
;;; result of the application of the reactions from the previous
;;; steps, plus the rest of the context sequence.  When the context
;;; sequence is empty, nothing is added to the current state.
(struct state (result rest-contexts) #:transparent)

;;; The dynamics of the reaction system only stores the reaction
;;; system itself.
(struct dynamics (rs) #:transparent
  #:methods gen:dds
  [;; Since reaction systems are deterministic, a singleton set is
   ;; produced, unless the context sequence is empty, in which case an
   ;; empty set of states is generated.  This transition is annotated
   ;; by the list of rules which were enabled in the current step.
   (define (dds-step-one-annotated dyn st)
     (define rs (dynamics-rs dyn))
     (define (apply-rs-annotate s rest-ctx)
       (define en (list-enabled rs s))
       (set (cons (list->set en)
                  (state (union-products rs en) rest-ctx))))
     (match st
       [(state res (cons ctx rest-ctx))
        (apply-rs-annotate (set-union res ctx) rest-ctx)]
       [(state res '()) (set)]))])

;;; Builds the state graph of a reaction system driven by a given
;;; context sequence.
(define (build-interactive-process-graph rs contexts)
  (dds-build-state-graph-annotated (dynamics rs)
                                   (set (state (set) contexts))))

;;; Builds the reduced state graph of a reaction system driven by
;;; a given context sequence.  Unlike build-interactive-process-graph,
;;; the nodes of this state graph do not contain the context sequence.
(define (build-reduced-state-graph rs contexts)
  (define sgr (build-interactive-process-graph rs contexts))
  (weighted-graph/directed
   (for/list ([e (in-edges sgr)])
     (define u (car e)) (define v (cadr e))
     (list (edge-weight sgr u v) (state-result u) (state-result v)))))

(module+ test
  (test-case "build-reduced-state-graph"
    (define rs (hash 'a (reaction (set 'x) (set 'y) (set 'z))
                     'b (reaction (set 'x) (set) (set 'y))))
    (define ctx (list (set 'x) (set 'y) (set 'z) (set) (set 'z)))
    (check-equal? (graphviz (build-reduced-state-graph rs ctx))
                  "digraph G {\n\tnode0 [label=\"(set)\\n\"];\n\tnode1 [label=\"(set 'y 'z)\\n\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node0 [label=\"#<set: #<set:>>\"];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"#<set: #<set: a b>>\"];\n\t\tnode1 -> node0 [label=\"#<set: #<set:>>\"];\n\t}\n}\n")))

(define (pretty-print-reduced-state-graph sgr)
  (update-graph sgr
                #:v-func (λ (st) (~a "{" (pretty-print-set st) "}"))
                #:e-func pretty-print-set-sets))

(module+ test
  (test-case "pretty-print-reduced-graph"
    (define rs (hash 'a (reaction (set 'x) (set 'y) (set 'z))
                     'b (reaction (set 'x) (set) (set 'y))))
    (define ctx (list (set 'x) (set 'y) (set 'z) (set) (set 'z)))
    (define sgr (build-reduced-state-graph rs ctx))
    (graphviz (pretty-print-reduced-state-graph sgr))))


;;; Builds the interactive process driven by the given context
;;; sequence.  The output is a list of pairs of lists in which the
;;; first element is the current context and the second element is the
;;; result of the application of reactions to the previous state.  The
;;; interactive process stops one step after the end of the context
;;; sequence, to show the effect of the last context.
(define (build-interactive-process rs contexts)
  (let ([dyn (dynamics rs)]
        [padded-contexts (append contexts (list (set)))])
    (for/fold ([proc '()]
               [st (state (set) padded-contexts)]
               #:result (reverse proc))
              ([c padded-contexts])
      (values
       (cons (match st
               [(state res ctx)
                (list (if (empty? ctx) (set) (car ctx)) res)])
             proc)
       (set-first (dds-step-one dyn st))))))

;;; Pretty-prints the context sequence and the current result of a
;;; state of the reaction system.  Note that we need to keep the full
;;; context sequence in the name of each state to avoid confusion
;;; between the states at different steps of the evolution.
(define/match (pretty-print-state st)
  [((state res ctx))
   (format "C:~a\nD:{~a}" (pretty-print-set-sets ctx) (pretty-print-set res))])

;;; Pretty prints the state graph of a reaction system.
(define (pretty-print-state-graph sgr)
  (update-graph sgr #:v-func pretty-print-state #:e-func pretty-print-set-sets))

(module+ test
  (test-case "Dynamics of reaction systems"
    (define r1 (reaction (set 'x) (set 'y) (set 'z)))
    (define r2 (reaction (set 'x) (set) (set 'y)))
    (define rs (make-immutable-hash (list (cons 'a r1) (cons 'b r2))))
    (define dyn (dynamics rs))
    (define state1 (state (set) (list (set 'x) (set 'y) (set 'z) (set) (set 'z))))
    (define sgr (dds-build-state-graph-annotated dyn (set state1)))
    (define ip (build-interactive-process-graph rs (list (set 'x) (set 'y) (set 'z) (set) (set 'z))))

    (check-equal? (dds-step-one-annotated dyn state1)
                  (set (cons
                        (set 'a 'b)
                        (state (set 'y 'z) (list (set 'y) (set 'z) (set) (set 'z))))))
    (check-equal? (dds-step-one dyn state1)
                  (set (state (set 'y 'z) (list (set 'y) (set 'z) (set) (set 'z)))))

    (check-true (has-vertex? sgr (state (set 'y 'z) (list (set 'y) (set 'z) (set) (set 'z)))))
    (check-true (has-vertex? sgr (state (set) (list (set 'z) (set) (set 'z)))))
    (check-true (has-vertex? sgr (state (set) (list (set) (set 'z)))))
    (check-true (has-vertex? sgr (state (set) (list (set 'z)))))
    (check-true (has-vertex? sgr (state (set) (list (set 'x) (set 'y) (set 'z) (set) (set 'z)))))
    (check-true (has-vertex? sgr (state (set) '())))

    (check-false (has-edge? sgr
                            (state (set) '())
                            (state (set) '())))
    (check-equal? (edge-weight sgr
                               (state (set 'y 'z) (list (set 'y) (set 'z) (set) (set 'z)))
                               (state (set) (list (set 'z) (set) (set 'z))))
                  (set (set)))
    (check-equal? (edge-weight sgr
                               (state (set) (list (set 'z) (set) (set 'z)))
                               (state (set) (list (set) (set 'z))))
                  (set (set)))
    (check-equal? (edge-weight sgr
                               (state (set) (list (set) (set 'z)))
                               (state (set) (list (set 'z))))
                  (set (set)))
    (check-equal? (edge-weight sgr
                               (state (set) (list (set 'z)))
                               (state (set) '()))
                  (set (set)))
    (check-equal? (edge-weight sgr
                               (state (set) (list (set 'x) (set 'y) (set 'z) (set) (set 'z)))
                               (state (set 'y 'z) (list (set 'y) (set 'z) (set) (set 'z))))
                  (set (set 'a 'b)))

    (check-equal? sgr ip)

    (check-equal? (build-interactive-process rs (list (set 'x) (set 'y) (set 'z) (set) (set 'z)))
                  (list
                   (list (set 'x) (set))
                   (list (set 'y) (set 'y 'z))
                   (list (set 'z) (set))
                   (list (set) (set))
                   (list (set 'z) (set))
                   (list (set) (set))))))