;; A beta distribution can be characterized either by parameters alpha and beta
;; or by its mean and standard deviation. These functions let us convert between
;; the two representations.

;; However there is a further complication: the value of sigma is bounded above by
;; sqrt(mu - mu^2). So the s parameter is going to be not sigma but a number between
;; 0 and 1 telling us how much of the bound we have. This isn't a problem for the functions
;; below other than to slightly adjust the outputs


(defun convert-to-us (ab-list)
   (let* 
        ((a (first ab-list))
         (b (second ab-list))
         (u (/ a (+ a b)))
         (s (sqrt (/ (* a b) (* (+ a b) (+ a b) (+ a b 1))))))
    
        (list u
              (/ s (sqrt (- u (* u u))))))
)

                        


(defun convert-to-ab (us-list)
   (let*
         ((u (first us-list))
          (s (* (second us-list) (sqrt (- u (* u u)))))
          (r (/ 1 (- 1 u)))
          (p (* (* r r) (* s s)))
          (b (/ (- (- r 1) p) (* r p))))
         (list (* (- r 1) b)
               b))
)


;; Set up some globals to hold the various parameters
;; We will just keep the us version around since that is what we control
;; Because of the above mentioned bound on sigma, the s parameter is really just
;; a multiplier -- it falls between 0 and 1 and tells us how close to the bound we
;; are. This means we have to be a little careful 

(def prior-params '(0.01 0.01))
(def posterior-params (convert-to-us '(1 1)))


;; Various vectors for the plots



(def x (rseq 0.02 0.98 49))
(def prior-y (apply 'beta-dens x (convert-to-ab prior-params)))
                      

;; The prior plot window

(def prior-plot (plot-lines x prior-y :title "Prior"))
(send prior-plot :range 1 0 3)


;; Stick in the new menu

(setf experiment-item
      (send menu-item-proto :new "Experiment"
                            :action 'get-heads-and-tails))

(setf quit-item
      (send menu-item-proto :new "Quit"
                            :action 'shut-down))

(send prior-plot :new-menu "Plot"
                 :items (list experiment-item
                              'dash
                              'rescale
                              'options
                              'dash
                              quit-item))


;; recompute and redraw the prior-plot 
(defun redraw-prior-plot ()
   (setf prior-y (apply 'beta-dens x (convert-to-ab prior-params)))
   (send prior-plot :clear nil)
   (send prior-plot :add-lines x prior-y)
  ; (send prior-plot :range 1 0 (ceiling (max prior-y)))
)


;; The two sliders to control the prior plot window

(def mean-slider (interval-slider-dialog (list 0.01 0.99)
                                         :points 99
                                         :title "Mean"
                        :action #'(lambda (p)
                                     (setf (select prior-params 0) p)
                                     (redraw-prior-plot)))
) 

(def spread-slider (interval-slider-dialog (list 1 99)
                                           :points 99
                                           :title "Spread %"
                        :action #'(lambda (p)
                                     (setf (select prior-params 1) (/ p 100))
                                     (redraw-prior-plot)))
) 



(defun shut-down ()
   (send spread-slider :close)
   (send mean-slider :close)
   (send prior-plot :close)
)










(def posterior-y (beta-dens x (+ 1 (first (convert-to-ab prior-params)))
                              (+ 1 (second (convert-to-ab prior-params)))))


;; We will just put the posterior plot on the same graph in another color -- this has
;; the advantage that if we move the sliders then it disappears (maybe not a good idea?)




(defun display-posterior (heads tails)
   (let*  
        ((posterior-y (beta-dens x (+ heads (first (convert-to-ab prior-params)))
                              (+ tails (second (convert-to-ab prior-params))))))
    

        (send prior-plot :add-lines x posterior-y :color 'magenta))
)






; Slightly modified version of p-dialog.lsp to give better parameter editing


;;;
;;; function to format a number or list to 4 digits
;;; used by edit-value-item-proto
;;; 

(defun fmt-list (x &optional (digits 3)) 
  (if (listp x) (format nil "(~{~,4a~#[~:; ~]~})"  x)
                (format nil "~,va" digits x)))
    
            

;;;
;;; edit-value-item-proto - for entering numeric data in a dialog
;;;

(defproto edit-value-item-proto () () edit-text-item-proto)

(defmeth edit-value-item-proto :value (&optional (value nil set))
  (if set
      (send self :text (format nil "~g" (fmt-list value)))
      (let ((txt (read (make-string-input-stream (send self :text)))))
        (if (listp txt)
            (if (numberp (first txt))
                txt
                (eval txt))
            (eval txt)))))

;;;
;;; parameter-item-proto - dialog item prototype for entering a parameter
;;;

(defproto parameter-item-proto '(label value-window))

;;; A new parameter item needs a name and initial value.  The length of the edit window
;;; and the size of the label associated with it may be specified.
;;; (The default edit-length of 36 is enough for a list of 5 numbers to fit comfortably.)
;;; If button is T, the label becomes a button that initiates the specified action.

(defmeth parameter-item-proto :isnew 
  (name value &key (edit-length 16) (label-size '(100 24)) (button nil) (action nil))
  (if button
      (setf (slot-value 'label) 
            (send button-item-proto :new name :action action :size label-size))
      (setf (slot-value 'label) 
            (send text-item-proto :new name :size label-size)))
  (setf (slot-value 'value-window) 
        (send edit-value-item-proto :new "" :text-length edit-length))
  (send (slot-value 'value-window) :value value))

;;; This is a convenience for putting a parameter item into a dialog.

(defmeth parameter-item-proto :dialog-item ()
  (list (slot-value 'label) (slot-value 'value-window)))

;;; Method for inserting or retrieving a value

(defmeth parameter-item-proto :value (&optional (value nil set))
  (if set (send (slot-value 'value-window) :value value)
      (send (slot-value 'value-window) :value)))

;;; Method to allow putting a messages in the window

(defmeth parameter-item-proto :text (text)
  (send (slot-value 'value-window) :text text))



;; Pop up a box to get the number of heads and tails
(defun get-heads-and-tails ()
  (let* (   
         (heads (send parameter-item-proto :new "Num Heads" 
                              1 
                              :edit-length 10))
         (tails (send parameter-item-proto :new "Num Tails" 
                              1 
                              :edit-length 10))
      
         (ok-button (send modal-button-proto :new "OK" :action
                          #'(lambda() (list 
                                            (send heads :value) 
                                            (send tails :value)
                                                             ))))
         (cancel-button (send modal-button-proto :new "Cancel"))
         (dialog (send modal-dialog-proto :new
                       (list  
                            (append (send heads :dialog-item)
                                    (send tails :dialog-item))
                            (list ok-button cancel-button))
                        :title "Experimental Outcome"           ))
         (answers (send dialog :modal-dialog)))
            (if answers
               (display-posterior (first answers) (second answers))))
)