;; lib/datalog/tests/aggregates.sx — count / sum / min / max.

(define dl-at-pass 0)
(define dl-at-fail 0)
(define dl-at-failures (list))

(define
  dl-at-deep=?
  (fn
    (a b)
    (cond
      ((and (list? a) (list? b))
        (and (= (len a) (len b)) (dl-at-deq-l? a b 0)))
      ((and (dict? a) (dict? b))
        (let ((ka (keys a)) (kb (keys b)))
          (and (= (len ka) (len kb)) (dl-at-deq-d? a b ka 0))))
      ((and (number? a) (number? b)) (= a b))
      (else (equal? a b)))))

(define
  dl-at-deq-l?
  (fn
    (a b i)
    (cond
      ((>= i (len a)) true)
      ((not (dl-at-deep=? (nth a i) (nth b i))) false)
      (else (dl-at-deq-l? a b (+ i 1))))))

(define
  dl-at-deq-d?
  (fn
    (a b ka i)
    (cond
      ((>= i (len ka)) true)
      ((let ((k (nth ka i)))
         (not (dl-at-deep=? (get a k) (get b k))))
        false)
      (else (dl-at-deq-d? a b ka (+ i 1))))))

(define
  dl-at-set=?
  (fn
    (a b)
    (and
      (= (len a) (len b))
      (dl-at-subset? a b)
      (dl-at-subset? b a))))

(define
  dl-at-subset?
  (fn
    (xs ys)
    (cond
      ((= (len xs) 0) true)
      ((not (dl-at-contains? ys (first xs))) false)
      (else (dl-at-subset? (rest xs) ys)))))

(define
  dl-at-contains?
  (fn
    (xs target)
    (cond
      ((= (len xs) 0) false)
      ((dl-at-deep=? (first xs) target) true)
      (else (dl-at-contains? (rest xs) target)))))

(define
  dl-at-test!
  (fn
    (name got expected)
    (if
      (dl-at-deep=? got expected)
      (set! dl-at-pass (+ dl-at-pass 1))
      (do
        (set! dl-at-fail (+ dl-at-fail 1))
        (append!
          dl-at-failures
          (str
            name
            "\n    expected: " expected
            "\n    got:      " got))))))

(define
  dl-at-test-set!
  (fn
    (name got expected)
    (if
      (dl-at-set=? got expected)
      (set! dl-at-pass (+ dl-at-pass 1))
      (do
        (set! dl-at-fail (+ dl-at-fail 1))
        (append!
          dl-at-failures
          (str
            name
            "\n    expected (set): " expected
            "\n    got:            " got))))))

(define
  dl-at-throws?
  (fn
    (thunk)
    (let
      ((threw false))
      (do
        (guard
          (e (#t (set! threw true)))
          (thunk))
        threw))))

(define
  dl-at-run-all!
  (fn
    ()
    (do
      ;; count
      (dl-at-test-set! "count siblings"
        (dl-query
          (dl-program
            "parent(p, bob). parent(p, alice). parent(p, charlie).
             sibling(X, Y) :- parent(P, X), parent(P, Y), !=(X, Y).
             sib_count(N) :- count(N, S, sibling(bob, S)).")
          (list (quote sib_count) (quote N)))
        (list {:N 2}))

      ;; sum
      (dl-at-test-set! "sum prices"
        (dl-query
          (dl-program
            "price(apple, 5). price(pear, 7). price(plum, 3).
             total(T) :- sum(T, X, price(F, X)).")
          (list (quote total) (quote T)))
        (list {:T 15}))

      ;; min
      (dl-at-test-set! "min score"
        (dl-query
          (dl-program
            "score(alice, 80). score(bob, 65). score(carol, 92).
             lo(M) :- min(M, S, score(P, S)).")
          (list (quote lo) (quote M)))
        (list {:M 65}))

      ;; max
      (dl-at-test-set! "max score"
        (dl-query
          (dl-program
            "score(alice, 80). score(bob, 65). score(carol, 92).
             hi(M) :- max(M, S, score(P, S)).")
          (list (quote hi) (quote M)))
        (list {:M 92}))

      ;; count over derived relation (stratification needed).
      (dl-at-test-set! "count over derived"
        (dl-query
          (dl-program
            "parent(a, b). parent(a, c). parent(b, d). parent(c, e).
             ancestor(X, Y) :- parent(X, Y).
             ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).
             num_ancestors(N) :- count(N, X, ancestor(a, X)).")
          (list (quote num_ancestors) (quote N)))
        (list {:N 4}))

      ;; count with no matches → 0.
      (dl-at-test-set! "count empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             zero(N) :- count(N, X, q(X)).")
          (list (quote zero) (quote N)))
        (list {:N 0}))

      ;; sum with no matches → 0.
      (dl-at-test-set! "sum empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             total(T) :- sum(T, X, q(X)).")
          (list (quote total) (quote T)))
        (list {:T 0}))

      ;; min with no matches → rule does not fire.
      (dl-at-test-set! "min empty"
        (dl-query
          (dl-program
            "p(1). p(2).
             lo(M) :- min(M, X, q(X)).")
          (list (quote lo) (quote M)))
        (list))

      ;; Aggregate with comparison filter on result.
      (dl-at-test-set! "popularity threshold"
        (dl-query
          (dl-program
            "post(p1). post(p2).
             liked(u1, p1). liked(u2, p1). liked(u3, p1).
             liked(u1, p2). liked(u2, p2).
             popular(P) :- post(P), count(N, U, liked(U, P)), >=(N, 3).")
          (list (quote popular) (quote P)))
        (list {:P (quote p1)}))

      ;; Aggregate vs single distinct.
      (dl-at-test-set! "distinct counted once"
        (dl-query
          (dl-program
            "rated(alice, x). rated(alice, y). rated(bob, x).
             rater_count(N) :- count(N, U, rated(U, F)).")
          (list (quote rater_count) (quote N)))
        (list {:N 2})))))

(define
  dl-aggregates-tests-run!
  (fn
    ()
    (do
      (set! dl-at-pass 0)
      (set! dl-at-fail 0)
      (set! dl-at-failures (list))
      (dl-at-run-all!)
      {:passed dl-at-pass
       :failed dl-at-fail
       :total (+ dl-at-pass dl-at-fail)
       :failures dl-at-failures})))