;; lib/prolog/tests/programs/nqueens.sx — N-queens via permute + safe.

(define pl-nq-test-count 0)
(define pl-nq-test-pass 0)
(define pl-nq-test-fail 0)
(define pl-nq-test-failures (list))

(define
  pl-nq-test!
  (fn
    (name got expected)
    (begin
      (set! pl-nq-test-count (+ pl-nq-test-count 1))
      (if
        (= got expected)
        (set! pl-nq-test-pass (+ pl-nq-test-pass 1))
        (begin
          (set! pl-nq-test-fail (+ pl-nq-test-fail 1))
          (append!
            pl-nq-test-failures
            (str name "\n    expected: " expected "\n    got:      " got)))))))

(define
  pl-nq-term-to-sx
  (fn
    (t)
    (cond
      ((pl-num? t) (pl-num-val t))
      ((pl-atom? t) (pl-atom-name t))
      (true (list :complex)))))

(define
  pl-nq-list-walked
  (fn
    (w)
    (cond
      ((and (pl-atom? w) (= (pl-atom-name w) "[]")) (list))
      ((and (pl-compound? w) (= (pl-fun w) ".") (= (len (pl-args w)) 2))
        (cons
          (pl-nq-term-to-sx (first (pl-args w)))
          (pl-nq-list-walked (nth (pl-args w) 1))))
      (true (list :not-list)))))

(define pl-nq-list-to-sx (fn (t) (pl-nq-list-walked (pl-walk-deep t))))

(define
  pl-nq-goal
  (fn
    (src env)
    (pl-instantiate (nth (first (pl-parse (str "g :- " src "."))) 2) env)))

(define
  pl-nq-prog-src
  "queens(L, Qs) :- permute(L, Qs), safe(Qs). permute([], []). permute(L, [H|T]) :- select(H, L, R), permute(R, T). select(X, [X|T], T). select(X, [H|T], [H|R]) :- select(X, T, R). safe([]). safe([Q|Qs]) :- safe(Qs), no_attack(Q, Qs, 1). no_attack(_, [], _). no_attack(Q, [Q1|Qs], D) :- is(D2, +(Q, D)), \\=(D2, Q1), is(D3, -(Q, D)), \\=(D3, Q1), is(D1, +(D, 1)), no_attack(Q, Qs, D1).")

(define pl-nq-db (pl-mk-db))
(pl-db-load! pl-nq-db (pl-parse pl-nq-prog-src))

(pl-nq-test!
  "queens([1], Qs) → 1 solution"
  (pl-solve-count!
    pl-nq-db
    (pl-nq-goal "queens([1], Qs)" {})
    (pl-mk-trail))
  1)

(pl-nq-test!
  "queens([1, 2], Qs) → 0 solutions"
  (pl-solve-count!
    pl-nq-db
    (pl-nq-goal "queens([1, 2], Qs)" {})
    (pl-mk-trail))
  0)

(pl-nq-test!
  "queens([1, 2, 3], Qs) → 0 solutions"
  (pl-solve-count!
    pl-nq-db
    (pl-nq-goal "queens([1, 2, 3], Qs)" {})
    (pl-mk-trail))
  0)

(pl-nq-test!
  "queens([1, 2, 3, 4], Qs) → 2 solutions"
  (pl-solve-count!
    pl-nq-db
    (pl-nq-goal "queens([1, 2, 3, 4], Qs)" {})
    (pl-mk-trail))
  2)

(pl-nq-test!
  "queens([1, 2, 3, 4, 5], Qs) → 10 solutions"
  (pl-solve-count!
    pl-nq-db
    (pl-nq-goal "queens([1, 2, 3, 4, 5], Qs)" {})
    (pl-mk-trail))
  10)

(define pl-nq-env-1 {})
(define pl-nq-goal-1 (pl-nq-goal "queens([1, 2, 3, 4], Qs)" pl-nq-env-1))
(pl-solve-once! pl-nq-db pl-nq-goal-1 (pl-mk-trail))

(pl-nq-test!
  "queens([1..4], Qs) first solution = [2, 4, 1, 3]"
  (pl-nq-list-to-sx (dict-get pl-nq-env-1 "Qs"))
  (list 2 4 1 3))

(define pl-nqueens-tests-run! (fn () {:failed pl-nq-test-fail :passed pl-nq-test-pass :total pl-nq-test-count :failures pl-nq-test-failures}))