From 28c545f31026103fdd4f399268548d81fabf871b Mon Sep 17 00:00:00 2001
From: Joshua Herring <jwherring@fastmail.com>
Date: Sun, 5 Apr 2026 21:51:56 -0400
Subject: [PATCH] generate enrollments

---
 main.go                                      | 60 ++++++++++++++++++--
 prompts/section_enrollment_data_generator.md | 17 ++++++
 sample_config.json                           |  1 +
 3 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 prompts/section_enrollment_data_generator.md

diff --git a/main.go b/main.go
index bc933fa..ef262d5 100644
--- a/main.go
+++ b/main.go
@@ -501,10 +501,11 @@ type Section struct {
 }
 
 type Config struct {
-	NumberOfUsers       int            `json:"number_of_users"`
-	NumberOfSections    int            `json:"number_of_sections"`
-	NumberOfInstructors int            `json:"number_of_instructors"`
-	PercentStatus       map[string]int `json:"percent_status"`
+	NumberOfUsers        int            `json:"number_of_users"`
+	NumberOfSections     int            `json:"number_of_sections"`
+	NumberOfInstructors  int            `json:"number_of_instructors"`
+	EnrollmentsPerStudent int           `json:"enrollments_per_student"`
+	PercentStatus        map[string]int `json:"percent_status"`
 }
 
 var statuses = []string{"JD", "LLM", "MCL", "PHD", "SJD", "MLS", "SPEC", "MNR", "CERT"}
@@ -626,6 +627,54 @@ func SaveGeneratedSections(sections []Section) {
 	fmt.Printf("Generated %d sections to sections.csv\n", len(sections))
 }
 
+type SectionEnrollment struct {
+	SectionID string `json:"section_id"`
+	UID       string `json:"uid"`
+}
+
+func GenerateSectionEnrollmentData(users []User, sections []Section, config Config) []SectionEnrollment {
+	if config.EnrollmentsPerStudent > len(sections) {
+		log.Fatalf("enrollments_per_student (%d) exceeds number of sections (%d)", config.EnrollmentsPerStudent, len(sections))
+	}
+
+	enrollments := make([]SectionEnrollment, 0, len(users)*config.EnrollmentsPerStudent)
+	for _, u := range users {
+		// Shuffle a copy of section indices and take the first EnrollmentsPerStudent
+		indices := rand.Perm(len(sections))
+		for _, idx := range indices[:config.EnrollmentsPerStudent] {
+			enrollments = append(enrollments, SectionEnrollment{
+				SectionID: sections[idx].SectionID,
+				UID:       u.UID,
+			})
+		}
+	}
+	return enrollments
+}
+
+func SaveGeneratedSectionEnrollments(enrollments []SectionEnrollment) {
+	f, err := os.Create("section_enrollments.csv")
+	if err != nil {
+		log.Fatalf("failed to create section_enrollments.csv: %v", err)
+	}
+	defer f.Close()
+
+	w := csv.NewWriter(f)
+	if err := w.Write([]string{"section_id", "uid"}); err != nil {
+		log.Fatalf("failed to write csv header: %v", err)
+	}
+	for _, e := range enrollments {
+		if err := w.Write([]string{e.SectionID, e.UID}); err != nil {
+			log.Fatalf("failed to write csv row: %v", err)
+		}
+	}
+	w.Flush()
+	if err := w.Error(); err != nil {
+		log.Fatalf("csv flush error: %v", err)
+	}
+
+	fmt.Printf("Generated %d section enrollments to section_enrollments.csv\n", len(enrollments))
+}
+
 func main() {
 	data, err := os.ReadFile("config.json")
 	if err != nil {
@@ -651,6 +700,9 @@ func main() {
 	generated_sections := GenerateSectionData(config)
 	SaveGeneratedSections(generated_sections)
 
+	generated_enrollments := GenerateSectionEnrollmentData(generated_users, generated_sections, config)
+	SaveGeneratedSectionEnrollments(generated_enrollments)
+
 }
 
 // uniqueNumericStrings generates n unique zero-padded numeric strings of the given digit length.
diff --git a/prompts/section_enrollment_data_generator.md b/prompts/section_enrollment_data_generator.md
new file mode 100644
index 0000000..9b387d4
--- /dev/null
+++ b/prompts/section_enrollment_data_generator.md
@@ -0,0 +1,17 @@
+Create a function GenerateSectionEnrollmentData in main.go that generates data for the following type:
+
+```go
+type SectionEnrollment struct {
+	SectionID string `json:"section_id"`
+	UID       string `json:"uid"`
+}
+```
+
+The function should take three arguments - users []User, sections []Section, config Config - and return a []SectionEnrollment
+
+It should generate a list of []SectionEnrollment equal in length to the length of the users list multiplied by the config.EnrollmentsPerStudent member meeting the following requirements:
+
+1. Each UID in the users list should be paird with config.EnrollmentsPerStudent distinct SectionIDs from the sections list 
+
+Create another function SaveGeneratedSectionEnrollments that takes a []SectionEnrollment, converts it into csv row data and saves it to a file section_enrollments.csv, creating this file if necessary and truncating any existing such file before writing.
+
diff --git a/sample_config.json b/sample_config.json
index 65ec688..12ba4bc 100644
--- a/sample_config.json
+++ b/sample_config.json
@@ -2,6 +2,7 @@
 	"number_of_users": 100,
 	"number_of_sections": 20,
 	"number_of_instructors": 8,
+	"enrollments_per_student": 4,
 	"percent_status": {
     "JD": 70,
     "LLM": 10,