#lang racket

;;; dds/utils

;;; Various utilities.

(require
 graph
 (for-syntax syntax/parse racket/list))

(provide
 ;; Functions
 (contract-out [eval-with (-> variable-mapping? any/c any)]
               [extract-symbols (-> any/c list?)]
               [any->string (-> any/c string?)]
               [stringify-variable-mapping (-> variable-mapping? string-variable-mapping?)]
               [string->any (-> string? any/c)]
               [read-org-table (-> string? (listof any/c))]
               [unstringify-pairs (-> (listof (general-pair/c string? any/c))
                                      (listof (general-pair/c symbol? any/c)))]
               [read-org-variable-mapping (-> string? variable-mapping?)]
               [update-vertices/unweighted (-> graph? (-> any/c any/c) graph?)]
               [update-graph (->* (graph?)
                                  (#:v-func (-> any/c any/c)
                                   #:e-func (-> any/c any/c))
                                  graph?)]
               [pretty-print-set (-> generic-set? string?)]
               [collect-by-key (-> (listof any/c) (listof any/c) (values (listof any/c) (listof (listof any/c))))]
               [collect-by-key/sets (-> (listof any/c) (listof any/c) (values (listof any/c) (listof (set/c any/c))))]

               [ht-values:list->set (-> (hash/c any/c (listof any/c)) (hash/c any/c (set/c any/c)))])
 ;; Contracts
 (contract-out [variable-mapping? contract?]
               [string-variable-mapping? contract?]
               [general-pair/c (-> contract? contract? contract?)])
 ;; Syntax
 auto-hash-ref/explicit auto-hash-ref/: sgfy unorg dotit)

;;; ===================
;;; HashTable Injection
;;; ===================

;;; This section of the file contains some utilities to streamline the
;;; usage of hash tables mapping symbols to values.  The goal is
;;; essentially to avoid having to write explicit hash-ref calls.

;;; A variable mapping is a hash table mapping symbols to values.
(define (variable-mapping? dict) (hash/c symbol? any/c))

;;; Given a (HashTable Symbol a) and a sequence of symbols, binds
;;; these symbols to the values they are associated to in the hash
;;; table, then puts the body in the context of these bindings.
;;;
;;; > (let ([ht #hash((a . 1) (b . 2))])
;;;     (auto-hash-ref/explicit (ht a b) (+ a (* 2 b))))
;;; 5
;;;
;;; Note that only one expression can be supplied in the body.
(define-syntax (auto-hash-ref/explicit stx)
  (syntax-parse stx
    [(_ (ht:id xs:id ...) body:expr)
     #`(let #,(for/list ([x (syntax->list #'(xs ...))])
                #`[#,x (hash-ref ht '#,x)])
         body)]))

;;; Given an expression and a (HashTable Symbol a), looks up the
;;; symbols with a leading semicolon and binds them to the value they
;;; are associated to in the hash table.
;;;
;;; > (let ([ht #hash((a . 1) (b . 2))])
;;;      (auto-hash-ref/: ht (+ :a (* 2 :b))))
;;; 5
;;;
;;; Note that the symbol :a is matched to the key 'a in the hash
;;; table.
;;;
;;; Note that only one expression can be supplied in the body.
(define-syntax (auto-hash-ref/: stx)
  (syntax-parse stx
    [(_ ht:id body)
     (let* ([names/: (collect-colons (syntax->datum #'body))])
       #`(let #,(for/list ([x names/:])
                  ;; put x in the same context as body
                  #`[#,(datum->syntax #'body x)
                     (hash-ref ht '#,(strip-colon x))])
           body))]))

;;; The helper functions for auto-hash-ref/:.
(begin-for-syntax
  ;; Collect all the symbols starting with a colon in datum.
  (define (collect-colons datum)
    (remove-duplicates
     (flatten
      (for/list ([token datum])
        (cond
          [(symbol? token)
           (let ([name (symbol->string token)])
             (if (eq? #\: (string-ref name 0))
                 token
                 '()))]
          [(list? token)
           (collect-colons token)]
          [else '()])))))

  ;; Strip the leading colon off x.
  (define (strip-colon x)
    (let ([x-str (symbol->string x)])
      (if (eq? #\: (string-ref x-str 0))
          (string->symbol (substring x-str 1))
          x))))

;;; Temporarily injects the mappings from the given hash table as
;;; bindings in a namespace including racket/base and then evaluates
;;; the expression.
;;;
;;; > (let ([ht #hash((a . 1) (b . 1))])
;;;     (eval-with ht '(+ b a 1)))
;;; 3
;;;
;;; The local bindings from the current lexical scope are not
;;; conserved.  Therefore, the following outputs an error about a
;;; missing identifier:
;;;
;;; > (let ([ht #hash((a . 1) (b . 1))]
;;;         [z 1])
;;;     (eval-with ht '(+ b z a 1)))
;;;
(define (eval-with ht expr)
  (parameterize ([current-namespace (make-base-namespace)])
    (for ([(x val) ht]) (namespace-set-variable-value! x val))
    (eval expr)))

;;; Same as eval-with, but returns only the first value produced by
;;; the evaluated expression.
(define (eval-with1 ht expr)
  (let ([vals (call-with-values (λ () (eval-with ht expr))
                                (λ vals vals))])
    (car vals)))


;;; ==============================
;;; Analysis of quoted expressions
;;; ==============================

;;; Produces a list of symbols appearing in the quoted expression
;;; passed in the first argument.
(define (extract-symbols form)
  (match form
    [(? symbol?) (list form)]
    [(? list?) (flatten (for/list ([x form])
                          (extract-symbols x)))]
    [else '()]))


;;; =========================
;;; Interaction with Org-mode
;;; =========================

;;; Org-mode supports laying out the output of code blocks as tables,
;;; which is very practical for various variable mappings (e.g.,
;;; states).  However, when the hash table maps variables to lists,
;;; Org-mode will create a column per list element, which may or may
;;; not be the desired effect.  In this section I define some
;;; utilities for nicer interoperation with Org-mode tables.  I also
;;; define some shortcuts to reduce the number of words to type when
;;; using dds with Org-mode.  See example.org for examples of usage.

;;; Converts any value to string.
(define (any->string x)
  (with-output-to-string (λ () (display x))))

;;; A string variable mapping is a mapping from variables to strings.
(define (string-variable-mapping? dict) (hash/c symbol? string?))

;;; Converts all the values of a variable mapping to string.
(define (stringify-variable-mapping ht)
  (for/hash ([(key val) ht]) (values key (any->string val))))

;;; A shortcut for variable-mapping-stingify.
(define-syntax-rule (sgfy ht) (stringify-variable-mapping ht))

;;; Reads any value from string.
(define (string->any str)
  (with-input-from-string str (λ () (read))))

;;; Reads a table from a string produced by Org-mode for a named
;;; table.  See example.org for examples.
(define (read-org-table str) (string->any str))

;;; A contract allowing pairs constructed via cons or via list.
(define (general-pair/c key-contract val-contract)
  (or/c (list/c key-contract val-contract)
        (cons/c key-contract val-contract)))

;;; Given a list of pairs of strings and some other values (possibly
;;; strings), converts the first element of each pair to a string, and
;;; reads the second element with string->any or keeps it as is if it
;;; is not a string.
(define (unstringify-pairs pairs)
  (for/list ([p pairs])
    (match p
      [(list key val)
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))]
      [(cons key val)                   ; also handle improper pairs
       (cons (string->symbol key) (if (string? val)
                                      (string->any val)
                                      val))])))

;;; Reads a variable mapping from a string, such as the one which
;;; Org-mode produces from tables.
(define (read-org-variable-mapping str)
  (make-hash (unstringify-pairs (read-org-table str))))

;;; A synonym for read-org-variable-mapping.
(define-syntax-rule (unorg str) (read-org-variable-mapping str))

;;; Typeset the graph via graphviz and display it.
(define-syntax-rule (dotit gr) (display (graphviz gr)))


;;; ==========================
;;; Additional graph utilities
;;; ==========================

;;; Apply a transformation to every vertex in the unweighted graph,
;;; return the new graph.  If the transformation function maps two
;;; vertices to the same values, these vertices will be merged in the
;;; resulting graph.  The transformation function may be called
;;; multiple times for the same vertex.
;;;
;;; This function does not rely on rename-vertex!, so it can be used
;;; to permute vertex labels.
(define (update-vertices/unweighted gr func)
  (unweighted-graph/directed
   (for/list ([e (in-edges gr)])
     (match-let ([(list u v) e])
       (list (func u) (func v))))))

;;; Given a graph, apply a transformation v-func to every vertex label
;;; and, if the graph is a weighted graph, the transformation e-func
;;; to every edge label.  Both transformations default to identity
;;; functions.  If gr is an weighted graph, the result is a weighted
;;; graph.  If gr is an unweighted graph, the result is an unweighted
;;; graph.
(define (update-graph gr #:v-func [v-func identity] #:e-func [e-func identity])
  (let ([edges (for/list ([e (in-edges gr)])
                 (match-let ([(list u v) e])
                   (if (unweighted-graph? gr)
                       (list (v-func u) (v-func v))
                       (list (e-func (edge-weight gr u v))
                             (v-func u) (v-func v)))))])
    (if (unweighted-graph? gr)
        (unweighted-graph/directed edges)
        (weighted-graph/directed edges))))


;;; ===============
;;; Pretty printing
;;; ===============

;;; Pretty print a set by listing its elements in alphabetic order.
(define (pretty-print-set s)
  (string-join (sort (set-map s any->string) string<?)))


;;; =========================
;;; Additional list utilities
;;; =========================

;;; Collects labels for duplicate edges into a sets of labels.
;;;
;;; More precisely, given a list of edges and weights, produces a new
;;; list of edges without duplicates, and a list of lists of weights
;;; in which each element corresponds to the edge (the input is
;;; suitable for graph constructors).
(define (collect-by-key edges labels)
  (for/fold ([ht (make-immutable-hash)]
             #:result (values (hash-keys ht) (hash-values ht)))
            ([e edges] [l labels])
    (hash-update ht e (λ (ls) (cons l ls)) empty)))

;;; Like collect-by-key, but returns a list of sets of weights.
(define (collect-by-key/sets edges labels)
  (let-values ([(es ls) (collect-by-key edges labels)])
    (values es (map list->set ls))))

;;; Converts the values of a hash table from lists to sets.
(define (ht-values:list->set ht)
  (for/hash ([(k v) (in-hash ht)])
    (values k (list->set v))))