#|

    class10.lsp     The EM algorithm 

    - Lowess simulation
    
    - Genetic example for EM, with likelihood calculation
    
    - Mixtures and the EM

|#


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                         ;;
;;    Smoothing simulation                                 ;;
;;                                                         ;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;--- Function to generate test data

(defun SINE-GEN (n sigma)
  (+ (sin (rseq 0 (* 2 pi) n))
     (* sigma (normal-rand n))
     ))

(def y (sine-gen 100 .5))
(def p (plot-points (iseq 100) y))
(send p :add-lines  (iseq 100)
      (setf truth 
	    (sin (rseq 0 (* 2 pi) 100))) :color 'magenta)

(send p :add-lines (setf fit (lowess (iseq 100) y)))
(setf fit (second fit))

(lowess '(1 2 3 4 5) '(2 4 3 1 2))


(def eplot (plot-lines (iseq 100) (- truth fit)))
(send eplot :add-lines (iseq 100) (repeat 0 100) :color 'magenta)

;--- Estimator

(defun MOVING-AVG (y width x)
  "Smooths assuming implicit x is equally spaced."
  ()
  )

(send eplot :range 1 -1.5 1)

;--- Simulator
(defun SIMULATE (stat nReps gen)
  ; From class 8
  (let ((results ())  )
    (dotimes (i nReps)
             (push (funcall stat (funcall gen))
                   results))
    results))

(setf simres (simulate #'(lambda (x)
			   (send eplot :add-lines
				 (iseq 100)
				 (- truth 
				    (second (lowess (iseq 100)
					 x :f .3)))))
		       200
                       #'(lambda () (sine-gen 100 .5))))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                         ;;
;;     Genetic example                                     ;;
;;                                                         ;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;--- genetic counts

  (def y '(125 18 20 34))


;--- here is the likelihood function g(Y)
  (defun log-g (p y)
    ; Ignores the multinomial constant
    (inner-product y
                   (log (list (+ .5 (/ p 4))    ; use some constants?
                              (/ (- 1 p) 4)
                              (/ (- 1 p) 4)
                              (/ p 4)))
                   ))
        
   (plot-function #'(lambda (p) (log-g p y)) .2 .8)


;--- Newton code from class 6

(defun FIXED-POINT (f x &optional (remDepth 10000))
  (format t "At ~6,3f~%" x)
  (let ((x1 (funcall f x))  )
    (if (< (abs (- x x1)) .0001)
        x1
        (if (< 0 remDepth) (fixed-point f x1 (1- remDepth)) x1)
        )))
(defun NUM-DERIV (f &key (eps .001))
  (lambda (x) (/ (- (funcall f (+ x eps))
                    (funcall f (- x eps)))
                 (* 2 eps)
                 )))
(defun NEWTON-EXT (f x0)
  "Args (f x0): Newton's method for finding min/max of f near x0."
  ; Needs some test to check whether its a max or a min.
  (let* ((fp  (num-deriv f))
         (fpp (num-deriv fp))   )
    (fixed-point #'(lambda (x) (- x (/ (funcall fp  x)
                                       (funcall fpp x)
                                       )))
                 x0)
    ))


;--- Try Newton's method  (efficient at one step)

(newton-ext #'(lambda (p) (log-g p y))
            .3)



;--- define some useful tools
  (defmacro WHILE (test &body body)
    `(do ()
         ((not ,test))
         ,@body))

  (defun NOT-CLOSE? (x y)
    (> (abs (- x y)) .0001))


;--- EM data and direct code

(def x '(x0 x1 18 20 34))

(defun ESTIMATE (p)
  ; Adjust the x counts
  (let ((sum (+ .5 (/ p 4)))  )
    (setf (select x 0) (* (/ .5 sum) (first y)))
    (setf (select x 1) (* (/ (/ p 4) sum) (first y)))
    )) 
  
(defun MAXIMIZE (x)
  ; compute new prob estimate
  (setf p^ (/ (+ (select x 1) (select y 3))
              (sum (select x 1) (select y '(1 2 3)))))
  )
 
;--- Do the calculation                 
(setf oldP 0)
(def p^ .3)
(while (not-close? oldP p^)
       (setf oldP p^)
       (estimate p^)
       (maximize x)
       (format t "P^ = ~5,3f with log-likelihood ~6,2f~%"
               p^ (log-g p^ y))
       )



;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                         ;;
;;     Normal mixtures example                             ;;
;;                           (fun with values)             ;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


;--- First, some code to generate the mixtures
   (defun MIXTURE-RAND (n m s)
     "Generate a sample of n observations from a normal mixture
      with parameters in the lists m and s"
     (let* ((k (length m))  ; and s
            (z (truncate (* k (uniform-rand n))))   )
       (setf trueZ z)
       (+ (select m z)
          (* (select s z) (normal-rand n))
          )))

(setf p (plot-lines 
         (kernel-dens 
          (setf y (mixture-rand 100 '(0 5) '(1 1)))) ) )

*plot-symbols*

(send p :add-points y (repeat 0.01 (length y)))
(send p :point-symbol (iseq (length y)) 
      (select '(cross square) truez))


;--- EM steps

(setf y  '(0 1 2 3 4 5))
(setf sd '(1 1))

   (defun ESTIMATE-Z-FROM-Y (y m s)
     ; Compute prob for each population, the E step
     (transpose
      (mapcar #'(lambda (p) (/ p (sum p)))
              (transpose
               (mapcar #'(lambda (mean sd)
                           (normal-dens (/ (- y mean) sd)))
                       m s))
              )))

(setf z^ (estimate-z-from-y y '(-2 2) sd))

   (defun ESTIMATE-M-FROM-DATA (y z)
     ; Compute new estimates of the mean and standard deviation
     (mapcar #'(lambda (wts)
                 (/ (sum (* y wts)) (sum wts)))
             z))

(setf parm^ (estimate-m-from-data y z^))

      
;--- Do some iterations

(def means '(-2 2))
(dotimes (i 5)
         (setf z^ (estimate-z-from-y y means sd))
         (setf means (estimate-m-from-data y z^))
         (format t "~d --- ~%" i)
         (format t "Means = ~{ ~7,3f ~}~%" means)
         )
         
(plot-points (first z^) trueZ :variable-labels '("Prob in First" "Group"))