#lang 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 dynamics% Dynamics% build-interactive-process-graph
 build-interactive-process-graph/simple-states
 pretty-print-state-graph/simple-states build-interactive-process
 build-interactive-process/org pretty-print-state pretty-print-state-graph
 )

(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))))
    (check-equal? (rs->ht-str-triples rs)
                  (hash 'a (list "x" "y" "z")
                        'b (list "x y" "" "t")))))

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

(define dynamics%
  (class (inst dds% State (Listof ReactionName))
    (super-new)
    (init-field [rs : ReactionSystem])
    (: step/annotated (-> State (Listof (Pairof (Listof ReactionName) State))))
    (define/override (step/annotated s)
      (match s
        [(state res (cons ctx rest-ctx))
         (define full-s (set-union ctx res))
         (define en (list-enabled rs full-s))
         (list (cons en (state (union-products rs en) rest-ctx)))]
        [(state _'()) '()]))))

(define-type Dynamics%
  (Instance (Class
             (init (rs ReactionSystem))
             (field (rs ReactionSystem))
             (build-state-graph (-> (Listof State) Graph))
             (build-state-graph*
              (-> (Listof State) (U 'full Exact-Positive-Integer) Graph))
             (build-state-graph*/annotated
              (-> (Listof State) (U 'full Exact-Positive-Integer) Graph))
             (build-state-graph/annotated (-> (Listof State) Graph))
             (step (-> State (Listof State)))
             (step* (-> (Listof State) (Listof State)))
             (step/annotated (-> State (Listof (Pairof (Listof Variable) State)))))))

(module+ test
  (test-case "dynamics%:step/annotated"
    (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                     'b (make-reaction '(x y) '() '(x))))
    (define dyn (new dynamics% [rs rs]))
    (define s0 (state (set 'x 'y)
                      (list (set) (set) (set 'x))))
    (define-values (_ 3-steps)
      (for/fold ([last-s : State s0]
                 [trace : (Listof (Pairof (Listof ReactionName) State)) '()])
                ([_ (in-range 1 4)])
        (define trans (send dyn step/annotated last-s))
        (values (cdar trans) (append trace trans))))
    (check-equal? 3-steps
                  (list
                   (cons '(b) (state (set 'x) (list (set) (set 'x))))
                   (cons '(a) (state (set 'z) (list (set 'x))))
                   (cons '(a) (state (set 'z) '()))))))

(: build-interactive-process-graph (-> ReactionSystem (Listof (Setof Species)) Graph))
(define (build-interactive-process-graph rs contexts)
  (send (new dynamics% [rs rs])
        build-state-graph/annotated
        (list (state (set) contexts))))

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

(: build-interactive-process-graph/simple-states (-> ReactionSystem (Listof (Setof Species)) Graph))
(define (build-interactive-process-graph/simple-states rs contexts)
  (define sgr (build-interactive-process-graph rs contexts))
  (weighted-graph/directed
   (for/list ([e (in-edges sgr)])
     (define u (assert-type (car e) State))
     (define v (assert-type (cadr e) State))
     (list (edge-weight sgr u v) (state-result u) (state-result v)))))

(module+ test
  (test-case "build-interactive-process-graph/simple-states"
    (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                     'b (make-reaction '(x y) '() '(x))))
    (define ctx : (Listof (Setof Species)) (list (set) (set) (set 'x)))
    (check-equal? (graphviz (build-interactive-process-graph/simple-states rs ctx))
                  "digraph G {\n\tnode0 [label=\"(set)\"];\n\tnode1 [label=\"(set 'z)\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node0 [label=\"'()\"];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"'(a)\"];\n\t}\n}\n")))

(: pretty-print-state-graph/simple-states (-> Graph Graph))
(define (pretty-print-state-graph/simple-states sgr)
  (update-graph
   sgr
   #:v-func
   (λ (st) (~a "{" (pretty-print-set (assert-type st (Setof Species))) "}"))
   #:e-func
   (λ (e) (pretty-print-set (assert-type e (Listof ReactionName))))))

(module+ test
  (test-case "pretty-print-state-graph/simple-states"
    (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                     'b (make-reaction '(x y) '() '(x))))
    (define ctx : (Listof (Setof Species)) (list (set) (set) (set 'x)))
    (check-equal?
     (graphviz (pretty-print-state-graph/simple-states
                (build-interactive-process-graph/simple-states rs ctx)))
     "digraph G {\n\tnode0 [label=\"{}\"];\n\tnode1 [label=\"{z}\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node0 [label=\"\"];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"a\"];\n\t}\n}\n")
    ))

(: build-interactive-process (-> ReactionSystem (Listof (Setof Species))
                                 (Listof (Pairof (Setof Species) (Setof Species)))))
(define (build-interactive-process rs contexts)
  (define dyn (new dynamics% [rs rs]))
  (define padded-contexts
    (append contexts (list (assert-type (set) (Setof Species)))))
  (for/fold ([proc : (Listof (Pairof (Setof Species) (Setof Species))) '()]
             [st : State (state (set) padded-contexts)]
             #:result (reverse proc))
            ([c padded-contexts])
    (define res (state-result st))
    (define ctx (state-rest-contexts st))
    (values
     ((inst cons (Pairof (Setof Species) (Setof Species)))
      (cons (if (empty? ctx) (assert-type (set) (Setof Species)) (car ctx)) res)
      proc)
     (set-first (send dyn step st)))))

(module+ test
  (test-case "build-interactive-process"
    (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                     'b (make-reaction '(x y) '() '(x))))
    (define ctx : (Listof (Setof Species)) (list (set 'x 'y) (set) (set 'x) (set)))
    (check-equal? (build-interactive-process rs ctx)
                  (list
                   (cons (set 'y 'x) (set))
                   (cons (set) (set 'x))
                   (cons (set 'x) (set 'z))
                   (cons (set) (set 'z))
                   (cons (set) (set))))))

(: build-interactive-process/org (-> ReactionSystem (Listof (Setof Species))
                                     (Listof (Listof (Setof Species)))))
(define (build-interactive-process/org rs context)
  (for/list : (Listof (Listof (Setof Species)))
            ([p (build-interactive-process rs context)])
    (list (car p) (cdr p))))

(module+ test
  (test-case "build-interactive-process/org"
    (define rs (hash 'a (make-reaction '(x) '(y) '(z))
                     'b (make-reaction '(x y) '() '(x))))
    (define ctx : (Listof (Setof Species)) (list (set 'x 'y) (set) (set 'x) (set)))
    (check-equal? (build-interactive-process/org rs ctx)
                  (list
                   (list (set 'y 'x) (set))
                   (list (set) (set 'x))
                   (list (set 'x) (set 'z))
                   (list (set) (set 'z))
                   (list (set) (set))))))

(: pretty-print-state (-> State String))
(define/match (pretty-print-state st)
  [((state res ctx))
   (format "C:~a\nD:{~a}" (pretty-print-set-sets ctx) (pretty-print-set res))])

(module+ test
  (test-case "pretty-print-state"
    (check-equal? (pretty-print-state
                   (state (set 'x 'y) (list (set 'z) (set) (set 'x))))
                  "C:{z}{}{x}\nD:{x y}")))

(: pretty-print-state-graph (-> Graph Graph))
(define (pretty-print-state-graph sgr)
  (update-graph
   sgr
   #:v-func (λ (st) (pretty-print-state (assert-type st State)))
   #:e-func (λ (e) (pretty-print-set (assert-type e (Listof ReactionName))))))

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