; 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 all list elements that satisfy a predicate.
(define (filter pred? list)
  (cond
    [(null? list) '()]
    [(pred? (car list)) (cons (car list) (filter pred? (cdr list)))]
    [else (filter pred? (cdr list))]))

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

; Four ways to make a process:
;   "end": stop with a final result (cf "reward" and "end" in decision.scm)
;   "say": produce some output then continue with the rest of the process
;     (cf "observe" and "do" in decision.scm and "discover" in influence.scm)
;   "ask": ask for some input then continue with a corresponding process
;     (cf "decide" and "switch" in decision.scm)
;   "pick": continue with a randomly picked process
;     (cf "observe" in decision.scm and "confuse" in influence.scm)
(define-struct end (result) (make-inspector))
(define-struct say (value then) (make-inspector))
(define-struct ask (choices) (make-inspector))
(define-struct pick (chances) (make-inspector))

; A choice is a branch off an "ask" process; it specifies a potential input and
; a corresponding process to continue with.
(define-struct choice (value then) (make-inspector))
(define (find-choice value choices)
  (find (lambda (choice) (equal? value (choice-value choice))) choices))

; A chance is a branch off a "pick" process; it specifies a probability
; distribution over processes to continue with.
(define-struct chance (probability then) (make-inspector))
(define (take-chance mass chances)
  (let ([remaining-mass (- mass (chance-probability (car chances)))])
    (if (positive? remaining-mass)
        (take-chance remaining-mass (cdr chances))
        (chance-then (car chances)))))

; It is useful to build "ask" and "pick" processes by applying the same
; continuation function to each of a list of values.
(define (ask values continuation)
  (make-ask (map (lambda (value)
                   (make-choice value (continuation value)))
                 values)))
(define (pick-uniform values continuation)
  (let ([probability (/ (length values))])
    (make-pick (map (lambda (value)
                      (make-chance probability (continuation value)))
                    values))))

; Three processes in the Monty Hall problem: the host and two contestants.
; The end result of the host is 1 or 0 to approximate the contestant's utility.
; The end result of contestant1 and contestant2 is just the symbol "done".
; The end result of contestant1a is the host's utility (that is, it is 1 if and
; only if the host follows the game-show rules and gets to keep the prize);
; otherwise contestant1a is just like contestant1.
(define doors '(a b c))
(define (neither x y) (lambda (z) (not (or (equal? x z) (equal? y z)))))
(define host
  (pick-uniform doors (lambda (prize)
    (ask doors (lambda (initial)
      (pick-uniform (filter (neither initial prize) doors) (lambda (open)
        (make-say open
          (ask doors (lambda (final)
            (make-say prize
              (make-end (if (equal? final prize) 1 0)))))))))))))
(define contestant1
  (pick-uniform doors (lambda (initial)
    (make-say initial
      (ask doors (lambda (open)
        (make-say initial
          (ask doors (lambda (prize)
            (make-end 'done))))))))))
(define contestant1a
  (pick-uniform doors (lambda (initial)
    (make-say initial
      (ask doors (lambda (open)
        (make-say initial
          (ask doors (lambda (prize)
            (make-end (if (or (equal? initial prize)
                              (equal? open prize)
                              (equal? initial open))
                          0 1)))))))))))
(define contestant2
  (pick-uniform doors (lambda (initial)
    (make-say initial
      (ask doors (lambda (open)
        (make-say (car (filter (neither initial open) doors))
          (ask doors (lambda (prize)
            (make-end 'done))))))))))

; The interpreter below simulates one run of a process using interactive user
; input and returns the end result.
(define (interact proc)
  (cond
    [(end? proc)
     (end-result proc)]
    [(say? proc)
     (printf "~a~n" (say-value proc))
     (interact (say-then proc))]
    [(ask? proc)
     (printf "~a? " (map choice-value (ask-choices proc)))
     (let ([choice (find-choice (read) (ask-choices proc))])
       (cond [choice (interact (choice-then choice))]
             [else (printf "That is not a valid choice value.~n")
                   (interact proc)]))]
    [(pick? proc)
     (interact (take-chance (random) (pick-chances proc)))]))

; The interpreter below simulates one run of two processes against each other
; and returns the end results.
(define (execute proc1 proc2)
  (cond
    [(pick? proc1)
     (execute (take-chance (random) (pick-chances proc1)) proc2)]
    [(pick? proc2)
     (execute proc1 (take-chance (random) (pick-chances proc2)))]
    [(end? proc1)
     (list (end-result proc1) (end-result proc2))]
    [(say? proc1)
     (printf "~a ->~n" (say-value proc1))
     (execute
       (say-then proc1)
       (choice-then (find-choice (say-value proc1) (ask-choices proc2))))]
    [(say? proc2)
     (printf "<- ~a~n" (say-value proc2))
     (execute
       (choice-then (find-choice (say-value proc2) (ask-choices proc1)))
       (say-then proc2))]))

; The rest of this file is concerned with solving a process, that is, finding a
; counter-process (if one exists) with the maximum expected utility.  The
; "solve" function solves a process and yields a solution, which consists of a
; counter-process (called the policy) and its expected utility.
(define-struct solution (policy utility) (make-inspector))
(define (solve proc)
  (solve-distr ((unpick 1 proc) '())))

(define (solve-distr distr)
  (cond
    [(end? (cdar distr))
     (make-solution (make-end 'done)
       (apply + (map (lambda (poss) (* (car poss) (end-result (cdr poss))))
                     distr)))]
    [(say? (cdar distr))
     (let ([collation
             (collate (lambda (value distr)
                        (let ([solution (solve-distr distr)])
                          (cons (make-choice value (solution-policy solution))
                                (solution-utility solution))))
                      (lambda (poss)
                        (cons (say-value (cdr poss))
                              (unpick (car poss) (say-then (cdr poss)))))
                      distr)])
       (make-solution (make-ask (map car collation))
                      (apply + (map cdr collation))))]
    [(ask? (cdar distr))
     (best-solution
       (map (lambda (choice)
              (let ([value (choice-value choice)])
                (let ([solution (solve-distr (conditionalize value distr '()))])
                  (make-solution (make-say value (solution-policy solution))
                                 (solution-utility solution)))))
            (ask-choices (cdar distr))))]))

(define (collate g f list)
  (let ([table (make-hash-table 'equal)])
    (for-each (lambda (element)
                (let ([pair (f element)])
                  (hash-table-put! table (car pair)
                    ((cdr pair) (hash-table-get table (car pair) '())))))
              list)
    (let ([collation '()])
      (hash-table-for-each table
        (lambda (k v) (set! collation (cons (g k v) collation))))
      collation)))

(define (conditionalize value distr distr0)
  (cond
    [(null? distr) distr0]
    [else ((unpick (caar distr)
                   (choice-then (find-choice value (ask-choices (cdar distr)))))
           (conditionalize value (cdr distr) distr0))]))

(define (unpick prob proc)
  (lambda (distr)
    (cond
      [(pick? proc)
       (unpick-chances prob (pick-chances proc) distr)]
      [else
       (cons (cons prob proc) distr)])))
(define (unpick-chances prob chances distr)
  (cond
    [(null? chances) distr]
    [else ((unpick (* prob (chance-probability (car chances)))
                   (chance-then (car chances)))
           (unpick-chances prob (cdr chances) distr))]))

(define (best-solution solutions)
  (let ([best-utility (apply max (map solution-utility solutions))])
    (find (lambda (solution) (= best-utility (solution-utility solution)))
          solutions)))