import pandas as pd
import numpy as np

np.random.seed(42)
n = 300

departments = np.random.choice(["Science", "Commerce", "Arts", "Engineering"], n)
attendance = np.random.choice(["Regular", "Irregular", "Borderline"], n, p=[0.4, 0.3, 0.3])
fav_colors = np.random.choice(["Red", "Blue", "Green", "Yellow", "Purple", "Orange"], n)
study_hours = np.round(np.random.uniform(2, 25, n), 1)
assign_sub = np.round(np.random.uniform(40, 100, n), 0)
backlogs = np.random.randint(0, 6, n)

# Stress logic
# High if: irregular attendance OR backlogs>=3 OR study_hours<5 OR assign_sub<55
stress = ((attendance == "Irregular") | 
          (backlogs >= 3) | 
          (study_hours < 5) | 
          (assign_sub < 55)).astype(int)

# Add some noise so it's not 100% predictable
flip = np.random.choice([0, 1], n, p=[0.9, 0.1])
stress = stress ^ flip  # XOR flip 10%

df = pd.DataFrame({
    "Student_ID": ["S{:03d}".format(i+1) for i in range(n)],
    "Department": departments,
    "Attendance_Level": attendance,
    "Favorite_Color": fav_colors,
    "Weekly_Study_Hours": study_hours,
    "Assignment_Submission_Pct": assign_sub.astype(int),
    "Backlogs_Count": backlogs,
    "Stress_Level": stress
})

df.to_csv("/root/fdp-workshop/fdp-colab-content/student_stress_data.csv", index=False)
print("Dataset generated: {} rows, {} cols".format(len(df), len(df.columns)))
print("Stress distribution:")
print(df["Stress_Level"].value_counts().to_string())
print("\nSample:")
print(df.head(3).to_string())
print("\nUnique colors: {}".format(df["Favorite_Color"].unique()))