#lang racket/base

(require
  racket/list
  "distribution.rkt")
(provide make-markov
         markov-add-example!
         markov-add-next-item!
         markov-generate
         markov->jsexpr
         jsexpr->markov)

; Markov chain structure
(struct markov
  (order ; order of the chain
   next-items)) ; table of given sequence -> distributions of next possible items

; Constructor
(define (make-markov [order 1])
  (markov order (make-hash)))

; Add a example to the generator. The example is a sequence of values in a list.
(define (markov-add-example! mrk example)
  (letrec
    ((fill
        (lambda (lst prev)
          (let ((distro (hash-ref (markov-next-items mrk) prev #f))
                (next (if (null? lst) #f (car lst))))
            (if distro
                (distribution-add-to! distro next)
                (let ((new-distro (make-distribution)))
                  (distribution-add-to! new-distro next)
                  (hash-set! (markov-next-items mrk) prev new-distro)))
            (when next (fill (cdr lst) (append (cdr prev) (list next))))
          ))))
    (fill example (make-list (markov-order mrk) #f))
  ))

; Add a possible next item from a sequence
; Use this to regenerate a markov chain from a serialized version (ie not from examples)
(define (markov-add-next-item! mrk seq next [occurences 1])
  (let ((distro (hash-ref (markov-next-items mrk) seq #f)))
    (if distro
        (distribution-add-to! distro next occurences)
        (let ((new-distro (make-distribution)))
          (distribution-add-to! new-distro next occurences)
          (hash-set! (markov-next-items mrk) seq new-distro)))))

; Generate a list with the given initial sequence
(define (generate-list next-items result previous)
  (let ((next (distribution-pick-from (hash-ref next-items previous))))
    (if (not next)
        (reverse result)
        (generate-list next-items (cons next result) (append (cdr previous) (list next))))))

; Get a suitable starting sequence from the given list
(define (find-starting-sequence next-items initial)
  ; If the given sequence is a suitable sequence, return it
  (if (hash-has-key? next-items initial)
      initial
      (letrec
        ; rating function to compare two sequences
        ((rate
            (lambda (result a b)
              (if (or (null? a) (null? b))
                  result
                  (rate (+ result (if (equal? (car a) (car b)) 1 0)) (cdr a) (cdr b))))))
        ; the suitable starting sequence is the one that get the best result with the rate function when compared to the initial sequence
        (caar
          (sort
            (map
              (lambda (lst) (cons lst (rate 0 (reverse lst) (reverse initial)))) ; The starting sequence is more suitable is we start from the last elements
              (shuffle
                (hash-keys next-items)))
            (lambda (a b) (> (cdr a) (cdr b))))))))

; Generate a list from the chain
(define (markov-generate mrk [initial #f])
  (if initial
      (let ((prev (find-starting-sequence (markov-next-items mrk) initial)))
        (generate-list (markov-next-items mrk) (reverse initial) prev))
      (generate-list (markov-next-items mrk) (list) (make-list (markov-order mrk) #f))))

; Serialisation to a json-compatible structure
; It is possible to pass a translator function for converting values that are not compatible to json
(define (markov->jsexpr mrk #:translator [tr (lambda (x) x)])
  (make-immutable-hash
    `((order . ,(markov-order mrk))
      (next . ,(hash-map
                 (markov-next-items mrk)
                 (lambda (k v)
                    (make-immutable-hash
                      `((sequence . ,(map tr k))
                        (distribution . ,(distribution->jsexpr v #:translator tr)))))))
      )))

; Deserialisation from a json-compatible structure
; It is possible to pass a translator function for converting values that are not compatible to json
(define (jsexpr->markov js #:translator [tr (lambda (x) x)])
  (define mrk (make-markov (hash-ref js 'order)))
  (for-each
    (lambda (x)
      (hash-set!
        (markov-next-items mrk)
        (map tr (hash-ref x 'sequence))
        (jsexpr->distribution (hash-ref x 'distribution) #:translator tr)))
    (hash-ref js 'next))
  mrk)