diff --git a/main.go b/main.go index d236e1a..e1b568e 100644 --- a/main.go +++ b/main.go @@ -10,12 +10,15 @@ import ( "encoding/json" "fmt" "html/template" - texttemplate "text/template" + "math" + "math/rand" "log" "net/http" "os" + "sort" "strconv" "strings" + texttemplate "text/template" _ "github.com/lib/pq" "google.golang.org/api/idtoken" @@ -77,6 +80,7 @@ func main() { http.HandleFunc("GET /api/trips/{tripID}/constraints", handleListConstraints(db)) http.HandleFunc("POST /api/trips/{tripID}/constraints", handleCreateConstraint(db)) http.HandleFunc("DELETE /api/trips/{tripID}/constraints/{constraintID}", handleDeleteConstraint(db)) + http.HandleFunc("POST /api/trips/{tripID}/solve", handleSolve(db)) http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { if err := db.Ping(); err != nil { http.Error(w, "db unhealthy", http.StatusServiceUnavailable) @@ -236,13 +240,13 @@ func handleListTrips(db *sql.DB) http.HandlerFunc { return } rows, err := db.Query(` - SELECT t.id, t.name, t.room_size, COALESCE( + SELECT t.id, t.name, t.room_size, t.prefer_not_multiple, COALESCE( json_agg(json_build_object('id', ta.id, 'email', ta.email)) FILTER (WHERE ta.id IS NOT NULL), '[]' ) FROM trips t LEFT JOIN trip_admins ta ON ta.trip_id = t.id - GROUP BY t.id, t.name, t.room_size + GROUP BY t.id, t.name, t.room_size, t.prefer_not_multiple ORDER BY t.id`) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) @@ -255,17 +259,18 @@ func handleListTrips(db *sql.DB) http.HandlerFunc { Email string `json:"email"` } type trip struct { - ID int64 `json:"id"` - Name string `json:"name"` - RoomSize int `json:"room_size"` - Admins []tripAdmin `json:"admins"` + ID int64 `json:"id"` + Name string `json:"name"` + RoomSize int `json:"room_size"` + PreferNotMultiple int `json:"prefer_not_multiple"` + Admins []tripAdmin `json:"admins"` } var trips []trip for rows.Next() { var t trip var adminsJSON string - if err := rows.Scan(&t.ID, &t.Name, &t.RoomSize, &adminsJSON); err != nil { + if err := rows.Scan(&t.ID, &t.Name, &t.RoomSize, &t.PreferNotMultiple, &adminsJSON); err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } @@ -384,14 +389,14 @@ func handleGetTrip(db *sql.DB) http.HandlerFunc { return } var name string - var roomSize int - err := db.QueryRow("SELECT name, room_size FROM trips WHERE id = $1", tripID).Scan(&name, &roomSize) + var roomSize, preferNotMultiple int + err := db.QueryRow("SELECT name, room_size, prefer_not_multiple FROM trips WHERE id = $1", tripID).Scan(&name, &roomSize, &preferNotMultiple) if err != nil { http.Error(w, "trip not found", http.StatusNotFound) return } w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]any{"id": tripID, "name": name, "room_size": roomSize}) + json.NewEncoder(w).Encode(map[string]any{"id": tripID, "name": name, "room_size": roomSize, "prefer_not_multiple": preferNotMultiple}) } } @@ -557,16 +562,36 @@ func handleUpdateTrip(db *sql.DB) http.HandlerFunc { return } var body struct { - RoomSize int `json:"room_size"` + RoomSize *int `json:"room_size"` + PreferNotMultiple *int `json:"prefer_not_multiple"` } - if err := json.NewDecoder(r.Body).Decode(&body); err != nil || body.RoomSize < 1 { - http.Error(w, "room_size must be at least 1", http.StatusBadRequest) + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, "invalid request body", http.StatusBadRequest) return } - _, err := db.Exec("UPDATE trips SET room_size = $1 WHERE id = $2", body.RoomSize, tripID) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return + if body.RoomSize != nil { + if *body.RoomSize < 1 { + http.Error(w, "room_size must be at least 1", http.StatusBadRequest) + return + } + } + if body.PreferNotMultiple != nil { + if *body.PreferNotMultiple < 1 { + http.Error(w, "prefer_not_multiple must be at least 1", http.StatusBadRequest) + return + } + } + if body.RoomSize != nil { + if _, err := db.Exec("UPDATE trips SET room_size = $1 WHERE id = $2", *body.RoomSize, tripID); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + } + if body.PreferNotMultiple != nil { + if _, err := db.Exec("UPDATE trips SET prefer_not_multiple = $1 WHERE id = $2", *body.PreferNotMultiple, tripID); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } } w.WriteHeader(http.StatusNoContent) } @@ -696,3 +721,310 @@ func handleDeleteConstraint(db *sql.DB) http.HandlerFunc { w.WriteHeader(http.StatusNoContent) } } + +func handleSolve(db *sql.DB) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + _, tripID, ok := requireTripAdmin(db, w, r) + if !ok { + return + } + + var roomSize, pnMultiple int + err := db.QueryRow("SELECT room_size, prefer_not_multiple FROM trips WHERE id = $1", tripID).Scan(&roomSize, &pnMultiple) + if err != nil { + http.Error(w, "trip not found", http.StatusNotFound) + return + } + + rows, err := db.Query("SELECT id, name FROM students WHERE trip_id = $1 ORDER BY id", tripID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + defer rows.Close() + var studentIDs []int64 + studentName := map[int64]string{} + for rows.Next() { + var id int64 + var name string + if err := rows.Scan(&id, &name); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + studentIDs = append(studentIDs, id) + studentName[id] = name + } + + if len(studentIDs) == 0 { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{"rooms": []any{}, "score": 0}) + return + } + + crows, err := db.Query(` + SELECT rc.student_a_id, rc.student_b_id, rc.kind::text, rc.level::text + FROM roommate_constraints rc + JOIN students sa ON sa.id = rc.student_a_id + WHERE sa.trip_id = $1`, tripID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + defer crows.Close() + + type constraint struct { + aID, bID int64 + kind, level string + } + var allConstraints []constraint + for crows.Next() { + var c constraint + if err := crows.Scan(&c.aID, &c.bID, &c.kind, &c.level); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + allConstraints = append(allConstraints, c) + } + + type pairKey struct{ a, b int64 } + overalls := map[pairKey]string{} + byPair := map[pairKey]map[string]string{} + for _, c := range allConstraints { + pk := pairKey{c.aID, c.bID} + if byPair[pk] == nil { + byPair[pk] = map[string]string{} + } + byPair[pk][c.level] = c.kind + } + levelPriority := []string{"admin", "parent", "student"} + for pk, levels := range byPair { + for _, lev := range levelPriority { + if kind, ok := levels[lev]; ok { + overalls[pk] = kind + break + } + } + } + + idx := map[int64]int{} + for i, id := range studentIDs { + idx[id] = i + } + n := len(studentIDs) + + mustTogether := map[[2]int]bool{} + mustApart := map[[2]int]bool{} + for pk, kind := range overalls { + ai, bi := idx[pk.a], idx[pk.b] + switch kind { + case "must": + p := [2]int{ai, bi} + if p[0] > p[1] { p[0], p[1] = p[1], p[0] } + mustTogether[p] = true + case "must_not": + p := [2]int{ai, bi} + if p[0] > p[1] { p[0], p[1] = p[1], p[0] } + mustApart[p] = true + } + } + + uf := make([]int, n) + for i := range uf { uf[i] = i } + var ufFind func(int) int + ufFind = func(x int) int { + if uf[x] != x { uf[x] = ufFind(uf[x]) } + return uf[x] + } + ufUnion := func(a, b int) { + ra, rb := ufFind(a), ufFind(b) + if ra != rb { uf[ra] = rb } + } + + for p := range mustTogether { + ufUnion(p[0], p[1]) + } + + hasConflict := false + for p := range mustApart { + if ufFind(p[0]) == ufFind(p[1]) { + hasConflict = true + break + } + } + + if hasConflict { + http.Error(w, "hard conflicts exist, resolve before solving", http.StatusBadRequest) + return + } + + groups := map[int][]int{} + for i := 0; i < n; i++ { + root := ufFind(i) + groups[root] = append(groups[root], i) + } + + score := func(assignment []int) int { + s := 0 + for pk, kind := range overalls { + ai, bi := idx[pk.a], idx[pk.b] + sameRoom := assignment[ai] == assignment[bi] + switch kind { + case "prefer": + if sameRoom { s++ } + case "prefer_not": + if sameRoom { s -= pnMultiple } + } + } + return s + } + + feasible := func(assignment []int) bool { + for p := range mustApart { + if assignment[p[0]] == assignment[p[1]] { return false } + } + roomCounts := map[int]int{} + for _, room := range assignment { + roomCounts[room]++ + } + for _, cnt := range roomCounts { + if cnt > roomSize { return false } + } + return true + } + + numRooms := (n + roomSize - 1) / roomSize + + assignment := make([]int, n) + groupList := make([][]int, 0, len(groups)) + for _, members := range groups { + groupList = append(groupList, members) + } + sort.Slice(groupList, func(i, j int) bool { + return len(groupList[i]) > len(groupList[j]) + }) + + roomCap := make([]int, numRooms) + for i := range roomCap { roomCap[i] = roomSize } + + placed := false + var placeGroups func(gi int) bool + placeGroups = func(gi int) bool { + if gi >= len(groupList) { return true } + grp := groupList[gi] + for room := 0; room < numRooms; room++ { + if roomCap[room] < len(grp) { continue } + ok := true + for _, member := range grp { + for p := range mustApart { + partner := -1 + if p[0] == member { partner = p[1] } + if p[1] == member { partner = p[0] } + if partner >= 0 && assignment[partner] == room { + alreadyPlaced := false + for gj := 0; gj < gi; gj++ { + for _, m := range groupList[gj] { + if m == partner { alreadyPlaced = true; break } + } + if alreadyPlaced { break } + } + if alreadyPlaced { ok = false; break } + } + } + if !ok { break } + } + if !ok { continue } + for _, member := range grp { assignment[member] = room } + roomCap[room] -= len(grp) + if placeGroups(gi + 1) { return true } + roomCap[room] += len(grp) + } + return false + } + placed = placeGroups(0) + + if !placed { + for i := 0; i < n; i++ { + assignment[i] = i % numRooms + } + } + + bestAssignment := make([]int, n) + copy(bestAssignment, assignment) + bestScore := score(assignment) + + temp := 10.0 + cooling := 0.9995 + iterations := 50000 + + for iter := 0; iter < iterations; iter++ { + i := rand.Intn(n) + j := rand.Intn(n) + if assignment[i] == assignment[j] { continue } + + iRoot := ufFind(i) + jRoot := ufFind(j) + + if iRoot == jRoot { continue } + + iGroup := groups[iRoot] + jGroup := groups[jRoot] + + iRoom := assignment[i] + jRoom := assignment[j] + + roomI := 0 + roomJ := 0 + for _, a := range assignment { + if a == iRoom { roomI++ } + if a == jRoom { roomJ++ } + } + + newRoomI := roomI - len(iGroup) + len(jGroup) + newRoomJ := roomJ - len(jGroup) + len(iGroup) + if newRoomI > roomSize || newRoomJ > roomSize { continue } + + for _, m := range iGroup { assignment[m] = jRoom } + for _, m := range jGroup { assignment[m] = iRoom } + + if !feasible(assignment) { + for _, m := range iGroup { assignment[m] = iRoom } + for _, m := range jGroup { assignment[m] = jRoom } + continue + } + + newScore := score(assignment) + delta := newScore - bestScore + if delta > 0 || rand.Float64() < math.Exp(float64(delta)/temp) { + if newScore > bestScore { + bestScore = newScore + copy(bestAssignment, assignment) + } + } else { + for _, m := range iGroup { assignment[m] = iRoom } + for _, m := range jGroup { assignment[m] = jRoom } + } + temp *= cooling + } + + type roomMember struct { + ID int64 `json:"id"` + Name string `json:"name"` + } + roomMap := map[int][]roomMember{} + for i, room := range bestAssignment { + sid := studentIDs[i] + roomMap[room] = append(roomMap[room], roomMember{ID: sid, Name: studentName[sid]}) + } + var rooms [][]roomMember + for room := 0; room < numRooms; room++ { + if members, ok := roomMap[room]; ok { + sort.Slice(members, func(i, j int) bool { return members[i].Name < members[j].Name }) + rooms = append(rooms, members) + } + } + sort.Slice(rooms, func(i, j int) bool { return rooms[i][0].Name < rooms[j][0].Name }) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{"rooms": rooms, "score": bestScore}) + } +} diff --git a/schema.sql b/schema.sql index 0fde698..665000b 100644 --- a/schema.sql +++ b/schema.sql @@ -10,11 +10,11 @@ END $$; CREATE TABLE IF NOT EXISTS trips ( id BIGSERIAL PRIMARY KEY, - name TEXT NOT NULL + name TEXT NOT NULL, + room_size INTEGER NOT NULL DEFAULT 2, + prefer_not_multiple INTEGER NOT NULL DEFAULT 5 ); -ALTER TABLE trips ADD COLUMN IF NOT EXISTS room_size INTEGER NOT NULL DEFAULT 2; - CREATE TABLE IF NOT EXISTS trip_admins ( id BIGSERIAL PRIMARY KEY, trip_id BIGINT NOT NULL REFERENCES trips(id) ON DELETE CASCADE, diff --git a/static/trip.html b/static/trip.html index fdf8028..c955e66 100644 --- a/static/trip.html +++ b/static/trip.html @@ -55,8 +55,9 @@ padding: 0 0.2rem; } .input-action:hover { opacity: 1; } - #trip-settings { display: flex; align-items: center; gap: 0.5rem; margin-bottom: 0.75rem; } - #room-size { width: 3.5rem; font-size: 0.85rem; padding: 0.2rem; border: 1px solid var(--wa-color-neutral-300, #ccc); border-radius: 0.25rem; } + #trip-settings { margin-bottom: 0.75rem; display: flex; flex-direction: column; gap: 0.3rem; } + #trip-settings label { display: flex; align-items: center; gap: 0.5rem; } + #trip-settings input[type="number"] { width: 3.5rem; font-size: 0.85rem; padding: 0.2rem; border: 1px solid var(--wa-color-neutral-300, #ccc); border-radius: 0.25rem; } .constraint-group { display: flex; flex-wrap: wrap; align-items: center; gap: 0.25rem; padding-bottom: 0.3rem; margin-bottom: 0.3rem; border-bottom: 1px solid #909090; } .constraint-level { font-size: 0.65rem; font-weight: bold; background: var(--wa-color-neutral-200, #ddd); color: var(--wa-color-neutral-700, #555); border-radius: 0.25rem; padding: 0.1rem 0.35rem; } .constraint-add { display: flex; gap: 0.5rem; align-items: center; margin-top: 0.3rem; } @@ -66,6 +67,11 @@ #hard-conflicts { margin-bottom: 0.5rem; } .conflict-row { margin-bottom: 0.2rem; } .conflict-icon { background: var(--wa-color-danger-50, #dc3545); color: white; border-radius: 0.15rem; padding: 0 0.15rem; font-size: 0.6rem; line-height: 1.2; vertical-align: middle; margin-right: 0.1rem; display: inline-block; } + #solver { margin-bottom: 0.75rem; } + #solver-results { margin-top: 0.5rem; } + .room-card { margin-bottom: 0.3rem; } + .room-label { font-weight: bold; font-size: 0.8rem; margin-bottom: 0.2rem; } + .solver-score { font-size: 0.8rem; margin-top: 0.3rem; color: var(--wa-color-neutral-500); }
@@ -79,12 +85,16 @@