generate enrollments

main.go

@@ -504,6 +504,7 @@ type Config struct {
 	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"`
 }
 
@@ -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.

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.
+

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,