test_data/generate_test_data.py

"""
Generate test data CSVs, sync_payload.json, and enroll_requests.json.

Run from project root:
    python test_data/generate_test_data.py

Outputs (written to test_data/):
  courses.csv          – 5 courses
  rooms.csv            – 3 rooms
  mentors.csv          – 9 mentors
  students.csv         – 200 students
  batches.csv          – 13 batches
  sync_payload.json    – POST to /api/v1/ingest/sync  (loads all entities)
  enroll_requests.json – batch_id → [student_ids] map; use enroll.py to load
  schedule_requests.json – sample AutoGenerateRequest per batch
"""

import csv
import json
import os

OUT = os.path.dirname(__file__)

# ── Master data ────────────────────────────────────────────────────────────────

COURSES = [
    {"id": "course-softskill",       "title": "Soft Skills",                    "type": "Addon",  "total_sessions": 4},
    {"id": "course-dsa",             "title": "Data Structures & Algorithms",   "type": "Major",  "total_sessions": 30},
    {"id": "course-data-analytics",  "title": "Data Analytics",                 "type": "Major",  "total_sessions": 15},
    {"id": "course-system-design",   "title": "System Design",                  "type": "Major",  "total_sessions": 10},
    {"id": "course-aiml",            "title": "AI / ML",                        "type": "Major",  "total_sessions": 12},
]

ROOMS = [
    {"id": "room-40", "name": "Room A", "capacity": 40},
    {"id": "room-50", "name": "Room B", "capacity": 50},
    {"id": "room-60", "name": "Room C", "capacity": 60},
]


def mentor_email(mentor_id):
    return f"{mentor_id}@example.com"

def mentor_mobile(mentor_id):
    return f"90000{mentor_id[-2:]}00"

MENTOR_COURSE_MAP = {
    "mentor-ss1": "course-softskill",
    "mentor-dsa1": "course-dsa",
    "mentor-dsa2": "course-dsa",
    "mentor-da1": "course-data-analytics",
    "mentor-da2": "course-data-analytics",
    "mentor-sd1": "course-system-design",
    "mentor-sd2": "course-system-design",
    "mentor-aiml1": "course-aiml",
    "mentor-aiml2": "course-aiml",
}

MENTORS = [
    {
        "id": "mentor-ss1", "name": "Priya Sharma",
        "email": mentor_email("mentor-ss1"),
        "mobile_number": mentor_mobile("mentor-ss1"),
        "expertise_tags": ["Soft Skills", "Communication", "Presentation"],
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat"],
        "course_id": MENTOR_COURSE_MAP["mentor-ss1"],
    },
    {
        "id": "mentor-dsa1", "name": "Rahul Kumar",
        "email": mentor_email("mentor-dsa1"),
        "mobile_number": mentor_mobile("mentor-dsa1"),
        "off_days": ["Sat", "Sun"],
        "course_id": MENTOR_COURSE_MAP["mentor-dsa1"],
    },
    {
        "id": "mentor-dsa2", "name": "Ankit Singh",
        "email": mentor_email("mentor-dsa2"),
        "mobile_number": mentor_mobile("mentor-dsa2"),
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri"],
        "course_id": MENTOR_COURSE_MAP["mentor-dsa2"],
    },
    {
        "id": "mentor-da1", "name": "Sneha Patel",
        "email": mentor_email("mentor-da1"),
        "mobile_number": mentor_mobile("mentor-da1"),
        "off_days": ["Sat", "Sun"],
        "course_id": MENTOR_COURSE_MAP["mentor-da1"],
    },
    {
        "id": "mentor-da2", "name": "Vijay Nair",
        "email": mentor_email("mentor-da2"),
        "mobile_number": mentor_mobile("mentor-da2"),
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri"],
        "course_id": MENTOR_COURSE_MAP["mentor-da2"],
    },
    {
        "id": "mentor-sd1", "name": "Kavya Reddy",
        "email": mentor_email("mentor-sd1"),
        "mobile_number": mentor_mobile("mentor-sd1"),
        "off_days": ["Sat", "Sun"],
        "course_id": MENTOR_COURSE_MAP["mentor-sd1"],
    },
    {
        "id": "mentor-sd2", "name": "Arjun Mehta",
        "email": mentor_email("mentor-sd2"),
        "mobile_number": mentor_mobile("mentor-sd2"),
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri"],
        "course_id": MENTOR_COURSE_MAP["mentor-sd2"],
    },
    {
        "id": "mentor-aiml1", "name": "Deepa Iyer",
        "email": mentor_email("mentor-aiml1"),
        "mobile_number": mentor_mobile("mentor-aiml1"),
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri"],
        "course_id": MENTOR_COURSE_MAP["mentor-aiml1"],
    },
    {
        "id": "mentor-aiml2", "name": "Rohan Gupta",
        "email": mentor_email("mentor-aiml2"),
        "mobile_number": mentor_mobile("mentor-aiml2"),
        "off_days": ["Mon", "Tue", "Wed", "Thu", "Fri"],
        "course_id": MENTOR_COURSE_MAP["mentor-aiml2"],
    },
]

# Batches – note: days/num_sessions are set at schedule-generation time, not stored here
BATCHES = [
    # Soft Skills – 1 mentor, 1 batch, Sunday only, 4 sessions
    {"id": "batch-ss-b1",      "name": "SS-B1 (Sun)",       "course_id": "course-softskill",      "mentor_id": "mentor-ss1",   "start_date": "2026-04-19", "status": "Active"},

    # DSA – mentor-dsa1: 2 weekday batches (Tue–Fri); mentor-dsa2: 2 weekend batches
    {"id": "batch-dsa-wd-b1",  "name": "DSA-WD-B1 (Tue-Fri)", "course_id": "course-dsa",           "mentor_id": "mentor-dsa1",  "start_date": "2026-04-14", "status": "Active"},
    {"id": "batch-dsa-wd-b2",  "name": "DSA-WD-B2 (Tue-Fri)", "course_id": "course-dsa",           "mentor_id": "mentor-dsa1",  "start_date": "2026-04-14", "status": "Active"},
    {"id": "batch-dsa-we-b1",  "name": "DSA-WE-B1 (Sat-Sun)", "course_id": "course-dsa",           "mentor_id": "mentor-dsa2",  "start_date": "2026-04-18", "status": "Active"},
    {"id": "batch-dsa-we-b2",  "name": "DSA-WE-B2 (Sat-Sun)", "course_id": "course-dsa",           "mentor_id": "mentor-dsa2",  "start_date": "2026-04-18", "status": "Active"},

    # Data Analytics – mentor-da1: 1 weekday batch; mentor-da2: 1 weekend batch
    {"id": "batch-da-wd-b1",   "name": "DA-WD-B1 (Mon-Fri)", "course_id": "course-data-analytics", "mentor_id": "mentor-da1",   "start_date": "2026-04-14", "status": "Active"},
    {"id": "batch-da-we-b1",   "name": "DA-WE-B1 (Sat-Sun)", "course_id": "course-data-analytics", "mentor_id": "mentor-da2",   "start_date": "2026-04-18", "status": "Active"},

    # System Design – mentor-sd1: 2 weekday batches; mentor-sd2: 2 weekend batches
    {"id": "batch-sd-wd-b1",   "name": "SD-WD-B1 (Mon-Fri)", "course_id": "course-system-design",  "mentor_id": "mentor-sd1",   "start_date": "2026-04-14", "status": "Active"},
    {"id": "batch-sd-wd-b2",   "name": "SD-WD-B2 (Mon-Fri)", "course_id": "course-system-design",  "mentor_id": "mentor-sd1",   "start_date": "2026-04-14", "status": "Active"},
    {"id": "batch-sd-we-b1",   "name": "SD-WE-B1 (Sat-Sun)", "course_id": "course-system-design",  "mentor_id": "mentor-sd2",   "start_date": "2026-04-18", "status": "Active"},
    {"id": "batch-sd-we-b2",   "name": "SD-WE-B2 (Sat-Sun)", "course_id": "course-system-design",  "mentor_id": "mentor-sd2",   "start_date": "2026-04-18", "status": "Active"},

    # AI/ML – 2 mentors, 1 weekend batch each
    {"id": "batch-aiml-we-b1", "name": "AIML-WE-B1 (Sat-Sun)", "course_id": "course-aiml",         "mentor_id": "mentor-aiml1", "start_date": "2026-04-18", "status": "Active"},
    {"id": "batch-aiml-we-b2", "name": "AIML-WE-B2 (Sat-Sun)", "course_id": "course-aiml",         "mentor_id": "mentor-aiml2", "start_date": "2026-04-18", "status": "Active"},
    # Attendance scenario batches – backdated/completed to test rich attendance states
    {"id": "batch-att-completed", "name": "ATT-Completed (backdated, DSA)", "course_id": "course-dsa",            "mentor_id": "mentor-dsa1", "start_date": "2026-01-06", "status": "Completed"},
    {"id": "batch-att-active",    "name": "ATT-Active (attendance history)", "course_id": "course-data-analytics", "mentor_id": "mentor-da1",  "start_date": "2026-03-01", "status": "Active"},
]

# ── Student generation ─────────────────────────────────────────────────────────
# 200 students spread across main course batches (one batch per student)
# DSA: 60 students (15 × 4 batches)
# DA:  40 students (20 × 2 batches)
# SD:  60 students (15 × 4 batches)
# AIML:40 students (20 × 2 batches)
# Total = 200

FIRST_NAMES = [
    "Aarav","Aditi","Akash","Ananya","Arjun","Avni","Chirag","Deepika","Divya","Gaurav",
    "Harsha","Isha","Jay","Kavya","Kiran","Lakshmi","Manish","Megha","Mohit","Neha",
    "Nikhil","Nisha","Pankaj","Pooja","Pranav","Priya","Rahul","Rajesh","Riya","Rohit",
    "Sachin","Sandeep","Sanjay","Sara","Shivam","Shreya","Simran","Sneha","Suresh","Tanvi",
    "Tarun","Usha","Varun","Vidya","Vikram","Vinay","Vivek","Yash","Zara","Amit",
]

LAST_NAMES = [
    "Sharma","Patel","Kumar","Singh","Gupta","Verma","Joshi","Mehta","Reddy","Nair",
    "Iyer","Agarwal","Mishra","Rao","Shah","Pillai","Bhat","Kaur","Pandey","Bajaj",
    "Jain","Malhotra","Khanna","Trivedi","Choudhary","Srivastava","Yadav","Tiwari","Das","Bose",
    "Sen","Mukherjee","Chatterjee","Banerjee","Ghosh","Roy","Dey","Chakraborty","Biswas","Paul",
]

def student_name(idx: int) -> str:
    first = FIRST_NAMES[idx % len(FIRST_NAMES)]
    last  = LAST_NAMES[(idx // len(FIRST_NAMES)) % len(LAST_NAMES)]
    return f"{first} {last}"

STUDENTS = [
    {
        "id": f"student-{i:03d}",
        "name": student_name(i - 1),
        "email": f"student-{i:03d}@example.com",
        "mobile_number": f"80000{i:03d}"
    }
    for i in range(1, 201)
] + [
    # Dedicated attendance-scenario students (stable IDs for test_api.py)
    {"id": "student-att-present", "name": "Test Present Student", "email": "student-att-present@example.com", "mobile_number": "70000001"},
    {"id": "student-att-late",    "name": "Test Late Student",    "email": "student-att-late@example.com",    "mobile_number": "70000002"},
    {"id": "student-att-absent",  "name": "Test Absent Student",  "email": "student-att-absent@example.com",  "mobile_number": "70000003"},
]

# ── Enrollment map ─────────────────────────────────────────────────────────────
# Maps batch_id → list of student ids enrolled in that batch
BATCH_ENROLLMENT: dict[str, list[str]] = {
    # DSA weekday
    "batch-dsa-wd-b1":  [f"student-{i:03d}" for i in range(1,  16)],   # 15 students
    "batch-dsa-wd-b2":  [f"student-{i:03d}" for i in range(16, 31)],   # 15 students
    # DSA weekend
    "batch-dsa-we-b1":  [f"student-{i:03d}" for i in range(31, 46)],   # 15 students
    "batch-dsa-we-b2":  [f"student-{i:03d}" for i in range(46, 61)],   # 15 students
    # DA weekday
    "batch-da-wd-b1":   [f"student-{i:03d}" for i in range(61,  81)],  # 20 students
    # DA weekend
    "batch-da-we-b1":   [f"student-{i:03d}" for i in range(81, 101)],  # 20 students
    # SD weekday
    "batch-sd-wd-b1":   [f"student-{i:03d}" for i in range(101, 116)], # 15 students
    "batch-sd-wd-b2":   [f"student-{i:03d}" for i in range(116, 131)], # 15 students
    # SD weekend
    "batch-sd-we-b1":   [f"student-{i:03d}" for i in range(131, 146)], # 15 students
    "batch-sd-we-b2":   [f"student-{i:03d}" for i in range(146, 161)], # 15 students
    # AIML weekend
    "batch-aiml-we-b1": [f"student-{i:03d}" for i in range(161, 181)], # 20 students
    "batch-aiml-we-b2": [f"student-{i:03d}" for i in range(181, 201)], # 20 students
    # Soft Skills addon – 20% of students (40 students, fits Room A capacity=40)
    # Picked from 4 different main batches (10 each)
    "batch-ss-b1": (
        [f"student-{i:03d}" for i in range(1,  11)] +   # 10 from DSA-WD-B1
        [f"student-{i:03d}" for i in range(61, 71)] +   # 10 from DA-WD-B1
        [f"student-{i:03d}" for i in range(101,111)] +  # 10 from SD-WD-B1
        [f"student-{i:03d}" for i in range(161,171)]    # 10 from AIML-WE-B1
    ),
    # Attendance scenario batches – 3 dedicated test students
    "batch-att-completed": ["student-att-present", "student-att-late", "student-att-absent"],
    "batch-att-active":    ["student-att-present", "student-att-late", "student-att-absent"],
}

# ── Schedule generation requests (sample for each batch) ──────────────────────
SCHEDULE_REQUESTS = [
    {"batch_id": "batch-ss-b1",      "start_date": "2026-04-19", "days": ["Sun"],                      "num_sessions": 4},
    {"batch_id": "batch-dsa-wd-b1",  "start_date": "2026-04-14", "days": ["Tue","Wed","Thu","Fri"],    "num_sessions": 30},
    {"batch_id": "batch-dsa-wd-b2",  "start_date": "2026-04-14", "days": ["Tue","Wed","Thu","Fri"],    "num_sessions": 30},
    {"batch_id": "batch-dsa-we-b1",  "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 30},
    {"batch_id": "batch-dsa-we-b2",  "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 30},
    {"batch_id": "batch-da-wd-b1",   "start_date": "2026-04-14", "days": ["Mon","Tue","Wed","Thu","Fri"], "num_sessions": 15},
    {"batch_id": "batch-da-we-b1",   "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 15},
    {"batch_id": "batch-sd-wd-b1",   "start_date": "2026-04-14", "days": ["Mon","Tue","Wed","Thu","Fri"], "num_sessions": 10},
    {"batch_id": "batch-sd-wd-b2",   "start_date": "2026-04-14", "days": ["Mon","Tue","Wed","Thu","Fri"], "num_sessions": 10},
    {"batch_id": "batch-sd-we-b1",   "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 10},
    {"batch_id": "batch-sd-we-b2",   "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 10},
    {"batch_id": "batch-aiml-we-b1", "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 12},
    {"batch_id": "batch-aiml-we-b2", "start_date": "2026-04-18", "days": ["Sat","Sun"],                "num_sessions": 12},
]

# ── Writers ────────────────────────────────────────────────────────────────────

def write_csv(filename: str, rows: list[dict], json_fields: list[str] | None = None):
    """Write a CSV. JSON fields are serialised as JSON strings."""
    path = os.path.join(OUT, filename)
    flat_rows = []
    for row in rows:
        flat = {}
        for k, v in row.items():
            if json_fields and k in json_fields:
                flat[k] = json.dumps(v)
            else:
                flat[k] = v
        flat_rows.append(flat)
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=flat_rows[0].keys())
        writer.writeheader()
        writer.writerows(flat_rows)
    print(f"  wrote {path}  ({len(rows)} rows)")


def write_json(filename: str, data):
    path = os.path.join(OUT, filename)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)
    print(f"  wrote {path}")


# ── Main ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("Generating test data ...")

    # CSVs (for upload/preview in the UI)
    write_csv("courses.csv",  COURSES)
    write_csv("rooms.csv",    ROOMS)
    write_csv("mentors.csv",  MENTORS,  json_fields=["expertise_tags", "off_days"])
    write_csv("mentors.csv",  MENTORS,  json_fields=["off_days"])
    write_csv("students.csv", STUDENTS)
    write_csv("batches.csv",  BATCHES)

    # sync_payload.json – POST body for /api/v1/ingest/sync
    sync_payload = {
        "courses":  COURSES,
        "rooms":    ROOMS,
        "mentors":  MENTORS,
        "students": STUDENTS,
        "batches":  BATCHES,
    }
    write_json("sync_payload.json", sync_payload)

    # enroll_requests.json – used by enroll.py
    write_json("enroll_requests.json", BATCH_ENROLLMENT)

    # schedule_requests.json – sample schedule generation requests
    write_json("schedule_requests.json", SCHEDULE_REQUESTS)

    print("\nDone.")
    print("\nNext steps:")
    print("  1. Load entities:  POST /api/v1/ingest/sync   with sync_payload.json")
    print("  2. Enroll students: python test_data/enroll.py --base-url http://localhost:7861")
    print("  3. Schedule:        POST /api/v1/schedule/auto-generate  (one per batch in schedule_requests.json)")