(in-package #:nekod.moe)

;; Sparse activation with noise injection (Noisy Top-K Gating)
;; Adds Gaussian noise during training to improve expert utilization

(defun gaussian-noise (mean stddev)
  "Box-Muller transform for Gaussian random number."
  (let* ((u1 (+ 1e-10 (random 1.0d0)))
         (u2 (random 1.0d0))
         (z (* (sqrt (* -2.0d0 (log u1))) (cos (* 2.0d0 pi u2)))))
    (+ mean (* stddev z))))

(defun noisy-top-k (router h &optional (noise-stddev 0.1d0) (training t))
  "Top-K gating with noise for training. In inference mode: no noise."
  (let* ((scores (all-scores router h))
         (noisy-scores
          (if training
              (mapcar (lambda (s)
                        (cons (car s)
                              (+ (cdr s) (gaussian-noise 0.0d0 noise-stddev))))
                      scores)
              scores))
         (sorted (sort noisy-scores #'> :key #'cdr))
         (selected (subseq sorted 0 (min (router-k router) (length sorted)))))
    (softmax-over-selected selected)))

;; Load balancing — prevent expert collapse
(defun load-balance-loss (router activations)
  "Auxiliary loss to prevent all traffic going to one expert.
   Returns scalar loss value — lower is more balanced."
  (let* ((n (length (router-expert-names router)))
         (counts (make-array n :initial-element 0)))
    (dolist (act activations)
      (let ((idx (position (car act) (router-expert-names router))))
        (when idx (incf (aref counts idx)))))
    (let ((total (reduce #'+ counts)))
      (if (zerop total) 0.0d0
          (loop for c across counts
                sum (let ((p (/ c total)))
                      (if (zerop p) 0.0d0 (* p (log p)))))))))