;;; wasp-markov --- Markov -*- lexical-binding: t; -*-
;;; Commentary:
;;; Code:

(require 'cl-lib)
(require 'rx)
(require 'dash)
(require 's)
(require 'ht)

(require 'wasp-utils)

(cl-defstruct (w/markov (:constructor w/markov-new))
  (words (ht-create)) ;; hashtable<word, hashtable<following word, weight>>
  )

(defun w/markov-feed (m s)
  "Feed string S to M."
  (let ((words (cons 'start (s-split (rx whitespace) s))))
    (while words
      (let ( (cur (car words))
             (next (cadr words)))
        (unless (ht-get (w/markov-words m) cur)
          (ht-set! (w/markov-words m) cur (ht-create)))
        (let ((freqs (ht-get (w/markov-words m) cur)))
          (ht-set! freqs next (+ 1 (ht-get freqs next 0))))
        (setf words (cdr words))))))

(defun w/markov-roll (ws)
  "Given weights WS, choose a random word."
  (-let
    [(total . weighted)
      (--reduce-from
        (let*
          ( (word (car it))
            (weight (cdr it))
            (new (+ (car acc) weight)))
          (cons
            new
            (cons
              (cons new word)
              (cdr acc))))
        (cons 0 nil)
        (ht->alist ws))]
    (-let [roll (random total)]
    (cdr (--first (< roll (car it)) (reverse weighted))))))

(defun w/markov-next (m prev)
  "Generate the next word for M given PREV."
  (when-let*
    ( (_ prev)
      (ws (ht-get (w/markov-words m) prev)))
    (w/markov-roll ws)))

(defun w/markov (m &optional fuel)
  "Generate a string from M. Optionally constrain to FUEL words."
  (let* ( (prev (w/markov-roll (ht-get (w/markov-words m) 'start)))
          (ret prev))
    (while-let
      ( (_ (> fuel 0))
        (next (w/markov-next m prev)))
      (setf ret (s-append (s-concat " " next) ret))
      (setf prev next)
      (cl-decf fuel))
    ret))

(defun w/markov-log (log)
  "Return a Markov chain trained on LOG."
  (let ( (ret (w/markov-new))
         (loglen (length log))
         (counter 0))
    (--each log
      (cl-incf counter)
      (message "processing %s/%s" counter loglen)
      (w/markov-feed ret (cdr it)))
    ret))

(provide 'wasp-markov)
;;; wasp-markov.el ends here