#lang racket/base

(require racket/list)
(require csp)

(define seating-chart (make-csp))

(define Npeople 6)

; return a series of seat variables from [start, end)
(define (chairs-subset start end) (make-var-names
                "seat"
                (range start end)))

; return all chairs
(define chairs
  (chairs-subset 0 Npeople))

(add-vars! seating-chart 
           chairs
           (range Npeople))

; constraint for people wanting to sit next to each other
; p1 wants to sit next to p2
(define (person-next-to p1 p2 a b c)
  (if (= b p1)
      (or (= a p2)
          (= c p2))
      #false))

; constraint for people avoiding each other
; p1 does not want to sit next to p2
(define (person-avoiding p1 p2 a b c)
  (if (= b p1)
      (and (not (= a p2))
          (not (= c p2)))
      #false))

; Constraints------------------------------------------------------------------

; # General constraints #######################################################
; 1. people can't occupy more than one seat
(add-pairwise-constraint! seating-chart
                          (lambda (a b) (not (= a b)))
                          chairs)

; # Solution-specific constraints (for a given problem) ######################
; These should be replaced with inputs, for a general user-submitted query

; person 2 wants to sit in seat 1, next to person 4
; TODO: this misses the ends (assuming the table is not just a line). person 2 in seat 0 and person 4 in seat 5 doesn't match, nor does person 2 in seat 5 and person 4 in seat 0
(add-constraint! seating-chart
                 (lambda (a b c) (person-next-to 2 4 a b c))
                 (chairs-subset 0 3))

; person 2 wants to sit next to person 4
(define (adjacent-constraint person1 person2)
  (map (add-constraint! seating-chart)))

; person 2 wants to sit in seat 1, but doesn't want to sit next to person 5
; TODO: this misses the ends (assuming the table is not just a line). person 2 in seat 0 and person 5 in seat 5 would incorrectly pass, as would person 2 in seat 5 and person 5 in seat 0
(add-constraint! seating-chart
                 (lambda (a b c) (person-avoiding 2 5 a b c))
                 (chairs-subset 0 3))

(printf "Size of state space: ~a.~n"
        (state-count seating-chart))

(time (solve* seating-chart 1))