;; lib/datalog/aggregates.sx — count / sum / min / max / findall.
;;
;; Surface form (always 3-arg after the relation name):
;;
;;   (count   Result  Var  GoalLit)
;;   (sum     Result  Var  GoalLit)
;;   (min     Result  Var  GoalLit)
;;   (max     Result  Var  GoalLit)
;;   (findall List    Var  GoalLit)
;;
;; Parsed naturally because arg-position compounds are already allowed
;; (Phase 4 needs them for arithmetic). At evaluation time the aggregator
;; runs `dl-find-bindings` on `GoalLit` under the current subst, collects
;; the distinct values of `Var`, and binds `Result`.
;;
;; Aggregation is non-monotonic — `count(C, X, p(X))` shrinks as p loses
;; tuples. The stratifier (lib/datalog/strata.sx) treats every aggregate's
;; goal relation as a negation-like edge so the inner relation is fully
;; derived before the aggregate fires.
;;
;; Empty input: count → 0, sum → 0, min/max → no binding (rule fails).

(define dl-aggregate-rels (list "count" "sum" "min" "max" "findall"))

(define
  dl-aggregate?
  (fn
    (lit)
    (and
      (list? lit)
      (>= (len lit) 4)
      (let ((rel (dl-rel-name lit)))
        (cond
          ((nil? rel) false)
          (else (dl-member-string? rel dl-aggregate-rels)))))))

;; Apply aggregation operator to a list of (already-distinct) numeric or
;; symbolic values. Returns the aggregated value, or :empty if min/max
;; has no input.
(define
  dl-do-aggregate
  (fn
    (op vals)
    (cond
      ((= op "count") (len vals))
      ((= op "sum") (dl-sum-vals vals 0))
      ((= op "findall") vals)
      ((= op "min")
        (cond
          ((= (len vals) 0) :empty)
          (else (dl-min-vals vals 1 (first vals)))))
      ((= op "max")
        (cond
          ((= (len vals) 0) :empty)
          (else (dl-max-vals vals 1 (first vals)))))
      (else (error (str "datalog: unknown aggregate " op))))))

(define
  dl-sum-vals
  (fn
    (vals acc)
    (cond
      ((= (len vals) 0) acc)
      (else (dl-sum-vals (rest vals) (+ acc (first vals)))))))

(define
  dl-min-vals
  (fn
    (vals i cur)
    (cond
      ((>= i (len vals)) cur)
      (else
        (let ((v (nth vals i)))
          (dl-min-vals vals (+ i 1) (if (< v cur) v cur)))))))

(define
  dl-max-vals
  (fn
    (vals i cur)
    (cond
      ((>= i (len vals)) cur)
      (else
        (let ((v (nth vals i)))
          (dl-max-vals vals (+ i 1) (if (> v cur) v cur)))))))

;; Membership check by deep equality (so 30 == 30.0 etc).
(define
  dl-val-member?
  (fn
    (v xs)
    (cond
      ((= (len xs) 0) false)
      ((dl-tuple-equal? v (first xs)) true)
      (else (dl-val-member? v (rest xs))))))

;; Evaluate an aggregate body lit under `subst`. Returns the list of
;; extended substitutions (0 or 1 element).
(define
  dl-eval-aggregate
  (fn
    (lit db subst)
    (let
      ((op (dl-rel-name lit))
        (result-var (nth lit 1))
        (agg-var (nth lit 2))
        (goal (nth lit 3)))
      (let
        ((substs (dl-find-bindings (list goal) db subst))
          (vals (list)))
        (do
          (for-each
            (fn
              (s)
              (let ((v (dl-apply-subst agg-var s)))
                (when
                  (not (dl-val-member? v vals))
                  (append! vals v))))
            substs)
          (let ((agg-val (dl-do-aggregate op vals)))
            (cond
              ((= agg-val :empty) (list))
              (else
                (let ((s2 (dl-unify result-var agg-val subst)))
                  (if (nil? s2) (list) (list s2)))))))))))

;; Stratification edges from aggregates: like negation, the goal's
;; relation must be in a strictly lower stratum so that the aggregate
;; fires only after the underlying tuples are settled.
(define
  dl-aggregate-dep-edge
  (fn
    (lit)
    (cond
      ((dl-aggregate? lit)
        (let ((goal (nth lit 3)))
          (cond
            ((and (list? goal) (> (len goal) 0))
              (let ((rel (dl-rel-name goal)))
                (if (nil? rel) nil {:rel rel :neg true})))
            (else nil))))
      (else nil))))