; For convenience, load tracing and pretty-printing functionality.
(require (lib "pretty.ss"))
(require (lib "trace"))

; A utility function to find a list element that satisfies a predicate.
(define (find pred? list)
  (cond
    [(null? list) #f]
    [(pred? (car list)) (car list)]
    [else (find pred? (cdr list))]))

; A utility function to find a list element that is "maximum" in some sense.
(define (max-by f list)
  (let ([m (apply max (map f list))])
    (find (lambda (element) (= m (f element))) list)))

; This line and (make-inspector) below tells DrScheme to
; show a visual representation of the contents of structures
(print-struct #t)

; A state has a name, an initial probability, a list of outgoing transition
; probabilities, and a list of observations.  Each observation is a list with
; two elements: a name and a probability.  The transition probabilities in each
; state and the observation probabilities in each state both add up to one.
(define-struct state (name initial transitions observations) (make-inspector))

; A hidden Markov model is a list of states.
(define hmm
  (list (make-state 'wake  .7 '(.8 .2) '((open .7) (closed .3)))
        (make-state 'sleep .3 '(.3 .7) '((open .2) (closed .8)))))

; Find the sequence of most likely states, given a sequence of observations.
; This computation uses forward and backward probabilities, defined below.
(define (path-best obss) (path-best-by + obss))
(define (path-best-by + obss)
  (map (lambda (f-probs b-probs)
         (max-by car (map (lambda (f-prob b-prob state)
                            (list (* f-prob b-prob) (state-name state)))
                          f-probs b-probs hmm)))
       (forward + obss) (backward + obss)))

; Find the most likely sequence of states, given a sequence of observations.
(define (best-path obss)
  (if (null? obss) '(1)
    (max-by car (map (lambda (state cont)
                       (cons (* (state-initial state) (car cont)) (cdr cont)))
                     hmm (best-path-obs (car obss) (cdr obss))))))
(define (best-path-obs obs obss)
  (map (lambda (state cont)
        (cons (* (cadr (assoc obs (state-observations state))) (car cont))
              (cons (state-name state) (cdr cont))))
       hmm (best-path-trans obss)))
(define (best-path-trans obss)
  (if (null? obss)
    (map (lambda (state) '(1)) hmm)
    (let ([continuations (best-path-obs (car obss) (cdr obss))])
      (map (lambda (state)
             (max-by car
               (map (lambda (trans cont) (cons (* trans (car cont)) (cdr cont)))
                    (state-transitions state) continuations)))
           hmm))))

; Compute forward probabilities.
(define (forward + obss)
  (if (null? obss) '()
    (forward-obs + (car obss) (cdr obss) (map state-initial hmm))))
(define (forward-obs + obs obss probs)
  (let ([probs (map (lambda (state prob)
                      (* prob (cadr (assoc obs (state-observations state)))))
                    hmm probs)])
    (cons probs (forward-trans + obss probs))))
(define (forward-trans + obss probs)
  (if (null? obss) '()
    (forward-obs + (car obss) (cdr obss)
      (apply map +
        (map (lambda (state prob)
               (map (lambda (trans) (* prob trans)) (state-transitions state)))
             hmm probs)))))

; Compute backward probabilities.
(define (backward + obss)
  (if (null? obss) '() (backward-trans + (cdr obss))))
(define (backward-obs + obs obss)
  (let ([probss (backward-trans + obss)])
    (cons (map (lambda (state prob)
                 (* (cadr (assoc obs (state-observations state))) prob))
               hmm (car probss))
          probss)))
(define (backward-trans + obss)
  (if (null? obss) (list (map (lambda (state) 1) hmm))
    (let ([probs-probss (backward-obs + (car obss) (cdr obss))])
      (cons (map (lambda (state) (apply + (map * (state-transitions state)
                                                 (car probs-probss))))
                 hmm)
            (cdr probs-probss)))))