ebisu.el

;; (defvar ebisu-tolerance 0.00001)
(defvar ebisu-tolerance 1e-8)
(defvar ebisu-max-iterations 100)

;; TODO: Can invert ratio to speed up computation
(defconst ebisu-phi (/ 2 (1+ (sqrt 5))))

;; TODO: Check code used by original Ebisu implementation
(defun ebisu-gss (fun a b)
  "Golden-section-search.
Basic variant without reuse of function evaluation.
https://en.wikipedia.org/wiki/Golden-section_search"
  (let ((c (- b (* (- b a) ebisu-phi)))
        (d (+ a (* (- b a) ebisu-phi)))
        (i 0))
    (while (and (< i ebisu-max-iterations)
                (> (abs (- b a)) ebisu-tolerance))
      (if (< (funcall fun c) (funcall fun d))
          (setq b d)
        (setq a c))
      ;; Recompute to avoid loss of precision
      (setq c (- b (* (- b a) ebisu-phi)))
      (setq d (+ a (* (- b a) ebisu-phi)))
      (incf i 1))
    (* (+ b a) 0.5)))

;; https://github.com/substack/gamma.js/blob/master/index.js
(defconst ebisu-g-ln (/ 607.0 128.0))
(defconst ebisu-p-ln
  '(0.99999999999999709182
    57.156235665862923517
    -59.597960355475491248
    14.136097974741747174
    -0.49191381609762019978
    0.33994649984811888699e-4
    0.46523628927048575665e-4
    -0.98374475304879564677e-4
    0.15808870322491248884e-3
    -0.21026444172410488319e-3
    0.21743961811521264320e-3
    -0.16431810653676389022e-3
    0.84418223983852743293e-4
    -0.26190838401581408670e-4
    0.36899182659531622704e-5))

(defun ebisu-gammaln (z)
  ;; JS Gamma returns NaN here
  (assert (>= z 0))
  (let ((x (car ebisu-p-ln)))
    (cl-loop for i from (1- (length ebisu-p-ln)) downto 1
             do (cl-incf x (/ (nth i ebisu-p-ln) (+ z i))))
    (let ((tt (+ z ebisu-g-ln 0.5)))
      (+
       ;; TODO: defconst
       (* 0.5 (log (* 2 pi)))
       (* (+ z 0.5) (log tt))
       (- tt)
       (log x)
       (- (log z))))))

(setq ebisu-gammaln--cache (make-hash-table))
(defun ebisu-gammaln-cached (z)
  (if-let ((res (gethash z ebisu-gammaln--cache)))
      res
    (let ((gamma (ebisu-gammaln z)))
      (puthash z gamma ebisu-gammaln--cache)
      gamma)))

;; Less accurate and possibly faster approximation
;; (defun ebisu-gammaln (x)
;;   "Lancoz Approximation of the natural logarithm of the gamma function.
;; Adapted from https://github.com/brown131/clmath/blob/master/functions/gamma.lisp"
;;   (if (< x 1.0)
;;       (let ((z (- 1.0 x)))
;;         (- (log (/ (* pi z) (sin (* pi z))))
;;            (ebisu-gammaln (+ 1.0 z))))
;;     (let* ((z   (- x 1.0))
;;            (tmp (+ z 5.0 0.5)))
;;       (+ (* (log tmp) (+ z 0.5))
;;          (- tmp)
;;          (log (sqrt (* 2 pi)))
;;          (log (+ 1.0
;;                  (/ 76.18009173 (+ z 1.0))
;;                  (/ -86.50532033 (+ z 2.0))
;;                  (/ 24.01409822 (+ z 3.0))
;;                  (/ -1.231739516 (+ z 4.0))
;;                  (/ 0.00120858003 (+ z 5.0))
;;                  (/ -0.00000536382 (+ z 6.0))))))))

;; TODO: Check if caching gammaln is useful here
(defun ebisu-betaln-ratio (a1 a b)
  (+ (- (ebisu-gammaln a1) (ebisu-gammaln (+ a1 b)))
     (- (ebisu-gammaln-cached (+ a b)) (ebisu-gammaln-cached a))))

;; TODO: Check if caching gammaln is useful here
(defun ebisu-betaln (a b)
  (- (+ (ebisu-gammaln-cached a)
        (ebisu-gammaln-cached b))
     (ebisu-gammaln-cached (+ a b))))

(defun ebisu-predict-recall (prior tnow)
  (cl-destructuring-bind (alpha beta t0) prior
   (let* ((dt (/ tnow t0))
          (ret (- (ebisu-betaln (+ alpha dt) beta)
                  (ebisu-betaln alpha beta)))
          ;; (ret (ebisu-betaln-ratio (+ alpha dt) alpha beta))
          )
     ;; NOTE: The python / js ebisu have a non-exact option
     ;; (if exact (exp ret) ret)
     ;; For use in org-fc we can probably default to exact results
     (exp ret))))

(cl-defun ebisu-default-model (tpri &optional (a 4.0) (b a))
  (list a b tpri))

(defun ebisu-binomln (n k)
  (- (+ (ebisu-betaln (- (1+ n) k) (1+ k))
        (log (1+ n)))))

(cl-defun ebisu-update-recall (prior successes total tnow &optional (rebalance t) (tback (caddr prior)))
  "Update recall model of PRIOR give SUCCESSES in TOTAL reviews at time TNOW.
REBALANCE defaults to true."
  (cl-destructuring-bind (alpha beta tpri) prior
    (let* ((tnow (float tnow))
           (tpri (float tpri))
           (tback (float tback))
           (dt (/ tnow tpri))
           (et (/ tback tnow))
           ;; TODO: This can be simplified if we assume only one review
           ;; per session
           (len (1+ (- total successes)))
           (binomlns
            (cl-loop for i below len
                     collecting (ebisu-binomln (- total successes) i)))
           ;; TODO: There is no reason to collect these into lists first
           (log-denominator
            (ebisu-logsumexp
             ;; TODO: Loop over binomlns, too
             (loop for i below len
                   collecting (+ (nth i binomlns)
                                 (ebisu-betaln beta
                                               (+ alpha
                                                  (* dt (+ successes i))
                                                  (* 0 dt et)))))
             ;; TODO: Handle this in optimibed logsumexp
             (loop for i below len collecting (if (evenp i) 1.0 -1.0))))
           (log-mean-num
            (ebisu-logsumexp
             (loop for i below len
                   collecting (+ (nth i binomlns)
                                 (ebisu-betaln beta
                                               (+ alpha
                                                  (* dt (+ successes i))
                                                  (* 1 dt et)))))
             (loop for i below len collecting (if (evenp i) 1.0 -1.0))))
           (log-m2-num
            (ebisu-logsumexp
             (loop for i below len
                   collecting (+ (nth i binomlns)
                                 (ebisu-betaln beta
                                               (+ alpha
                                                  (* dt (+ successes i))
                                                  (* 2 dt et)))))
             (loop for i below len collecting (if (evenp i) 1.0 -1.0))))
           (mean (exp (- log-mean-num log-denominator)))
           (m2 (exp (- log-m2-num log-denominator)))
           (mean-sq (exp (* 2 (- log-mean-num log-denominator))))
           (sig2 (- m2 mean-sq)))
      ;; TODO: Check if mean, m2, sig2 are finite & >= 0
      ;; NOTE: This inlines the `_meanVarToBeta` function found in the JS version
      (let* ((tmp (- (/ (* mean (- 1 mean)) sig2) 1))
             (alpha-new (* mean tmp))
             (beta-new (* (- 1 mean) tmp))
             (proposed (list alpha-new beta-new tback)))
        (if rebalance
            (ebisu-rebalance prior successes total tnow proposed)
          proposed)))))

(defun ebisu-rebalance (prior k n tnow proposed)
  (cl-destructuring-bind (alpha-new beta-new _t) proposed
    (if (or (> alpha-new (* 2 beta-new))
            (> beta-new (* 2 alpha-new)))
        (ebisu-update-recall
         prior k n tnow
         ;; no rebalancing
         nil
         ;; NOTE: Inlined JS variable `roughHalfLife`
         (ebisu-model-to-percentile-decay proposed 0.5 t))
      proposed)))

;; TODO: The JS version also accepts a `tolerance` argument
(cl-defun ebisu-model-to-percentile-decay (model &optional (percentile 0.5) coarse)
  (assert (<= 0 percentile 1))
  (cl-destructuring-bind (alpha beta t0) model
    (let* ((log-bab (ebisu-betaln alpha beta))
           (log-percentile (log percentile))
           (f (lambda (lndelta)
                (-
                 ;; NOTE: Inlined `logMean` of the JS version
                 (ebisu-betaln (+ alpha (exp lndelta)) beta)
                 log-bab
                 log-percentile)))
           (bracket-width (if coarse 1 6))
           (b-low (* bracket-width -0.5))
           (b-high (* bracket-width 0.5))
           (f-low (funcall f b-low))
           (f-high (funcall f b-high)))

      (while (and (> f-low 0) (> f-high 0))
        ;; Move bracket up
        (setq b-low b-high)
        (setq f-low f-high)
        (setq b-high (+ b-high bracket-width))
        (setq f-high (funcall f b-high)))

      (while (and (< f-low 0) (< f-high 0))
        ;; Move bracket down
        (setq b-high b-low)
        (setq f-high f-low)
        (setq b-low (- b-low bracket-width))
        (setq f-low (funcall f b-low)))

      ;; Failed to bracket
      ;; TODO: Meaningful error messages
      (assert (and (> f-low 0) (< f-high 0)))

      ;; TODO: The JS version handles `coarse = true` here
      ;; TODO: The JS version passes the tolerance parameter to gss
      ;; (called `minimize-golden-seciton-1d`)
      (if coarse
          (* (+ (exp b-low) (exp b-high)) 0.5 t0)
        ;; TODO: Check if gss converged
        (* (exp (ebisu-gss
                 (lambda (x) (abs (funcall f x)))
                 b-low b-high))
           t0)))))

(defun ebisu-logsumexp (as bs)
  (let* ((a-max (apply #'max as))
         ;; TODO: In the JS version, the loop runs backwards,
         ;; maybe that has something to do with numerical stability?
         (s (cl-loop for a in as
                     for b in bs
                     summing (* b (exp (- a a-max))))))
    ;; NOTE: The JS version also returns a sign but that is never used
    (+ (log (abs s)) a-max)))