server/data_generator.py

import pandas as pd
from datetime import datetime, timedelta

# ---------------------------------------------------------------------------
# Cluster A: filter → sort → extract
# ---------------------------------------------------------------------------

_A1_df = pd.DataFrame({
    "product": ["Widget", "Gadget", "Doohickey", "Sprocket", "Thingamajig", "Gizmo"],
    "revenue": [5000, 8000, 3000, 7500, 2000, 6000],
    "region": ["West", "West", "East", "West", "East", "North"],
})

_A2_df = pd.DataFrame({
    "name": ["Alice", "Bob", "Carol", "Dan", "Eve", "Frank", "Grace"],
    "dept": ["Eng", "Sales", "Eng", "Eng", "HR", "Eng", "Sales"],
    "tenure": [5, 3, 8, 2, 4, 10, 1],
})

_A3_df = pd.DataFrame({
    "order_id": ["ORD-101", "ORD-102", "ORD-103", "ORD-104", "ORD-105", "ORD-106"],
    "status": ["shipped", "pending", "shipped", "shipped", "cancelled", "shipped"],
    "quantity": [50, 20, 80, 30, 10, 60],
})

_A4_df = pd.DataFrame({
    "student_id": ["S01", "S02", "S03", "S04", "S05", "S06", "S07"],
    "grade": ["A", "B", "A", "A", "C", "A", "B"],
    "gpa": [3.9, 3.2, 3.95, 3.7, 2.8, 3.85, 3.1],
})

_A5_df = pd.DataFrame({
    "item_name": ["Bolts", "Nails", "Screws", "Washers", "Rivets", "Pins"],
    "stock": [5, 50, 8, 3, 100, 7],
    "reorder_priority": [2, 10, 3, 1, 15, 4],
})

# ---------------------------------------------------------------------------
# Cluster B: normalize → filter → extract
# ---------------------------------------------------------------------------

_B1_df = pd.DataFrame({
    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6"],
    "email": ["Alice@Gmail.COM", "bob@yahoo.com", "Carol@GMAIL.com",
              "dan@gmail.COM", "eve@outlook.com", "frank@Gmail.com"],
})

_B2_df = pd.DataFrame({
    "id": [1, 2, 3, 4, 5, 6, 7],
    "name": ["  Alice  ", " Bob", "  Andrew ", "Anna  ", " Carl ", "  Amy", " Brian  "],
})

_B3_df = pd.DataFrame({
    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6"],
    "country_code": ["us", "UK", "Us", "ca", "US", "us"],
})

_B4_df = pd.DataFrame({
    "product_id": ["P1", "P2", "P3", "P4", "P5", "P6"],
    "product_name": ["widget pro", "basic gadget", "pro sprocket",
                     "mega pro tool", "simple bolt", "pro widget x"],
})

_B5_df = pd.DataFrame({
    "contact_id": ["C1", "C2", "C3", "C4", "C5", "C6"],
    "phone": ["(555) 123-4567", "555.987.6543", "(555)111-2222",
              "5551234", "555-999-8888", "(555) 000 1111"],
})

# ---------------------------------------------------------------------------
# Cluster C: date delta → filter → count
# ---------------------------------------------------------------------------

_today = datetime(2026, 3, 8)

_C1_df = pd.DataFrame({
    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6", "U7"],
    "signup_date": [
        _today - timedelta(days=10),
        _today - timedelta(days=45),
        _today - timedelta(days=5),
        _today - timedelta(days=90),
        _today - timedelta(days=25),
        _today - timedelta(days=3),
        _today - timedelta(days=60),
    ],
    "active": [True, True, True, False, True, True, False],
})

_C2_df = pd.DataFrame({
    "order_id": ["O1", "O2", "O3", "O4", "O5", "O6"],
    "order_date": [
        _today - timedelta(days=2),
        _today - timedelta(days=10),
        _today - timedelta(days=5),
        _today - timedelta(days=1),
        _today - timedelta(days=14),
        _today - timedelta(days=6),
    ],
    "amount": [100, 200, 150, 50, 300, 75],
})

_C3_df = pd.DataFrame({
    "emp_id": ["E1", "E2", "E3", "E4", "E5", "E6"],
    "hire_date": [
        _today - timedelta(days=365),
        _today - timedelta(days=900),
        _today - timedelta(days=200),
        _today - timedelta(days=500),
        _today - timedelta(days=100),
        _today - timedelta(days=1500),
    ],
    "dept": ["Eng", "Sales", "Eng", "HR", "Eng", "Sales"],
})

_C4_df = pd.DataFrame({
    "event_id": ["EV1", "EV2", "EV3", "EV4", "EV5", "EV6"],
    "event_date": [
        _today + timedelta(days=5),
        _today + timedelta(days=20),
        _today + timedelta(days=10),
        _today + timedelta(days=3),
        _today + timedelta(days=30),
        _today + timedelta(days=12),
    ],
    "venue": ["Hall A", "Hall B", "Hall A", "Hall C", "Hall B", "Hall A"],
})

_C5_df = pd.DataFrame({
    "user_id": ["U1", "U2", "U3", "U4", "U5", "U6", "U7", "U8"],
    "birthdate": [
        _today - timedelta(days=365 * 20),
        _today - timedelta(days=365 * 30),
        _today - timedelta(days=365 * 22),
        _today - timedelta(days=365 * 17),
        _today - timedelta(days=365 * 19),
        _today - timedelta(days=365 * 25),
        _today - timedelta(days=365 * 24),
        _today - timedelta(days=365 * 15),
    ],
})

# ---------------------------------------------------------------------------
# Cluster D: groupby → aggregate → sort → slice
# ---------------------------------------------------------------------------

_D1_df = pd.DataFrame({
    "region": ["West", "East", "West", "North", "East", "North", "West", "East", "South", "South"],
    "revenue": [5000, 3000, 7000, 4000, 6000, 2000, 3000, 8000, 1000, 5000],
    "product": ["A", "B", "C", "A", "B", "C", "A", "B", "C", "A"],
})

_D2_df = pd.DataFrame({
    "customer": ["Alice", "Bob", "Alice", "Carol", "Bob", "Alice",
                 "Carol", "Bob", "Alice", "Bob", "Carol", "Bob",
                 "Alice", "Bob"],
    "order_id": [f"O{i}" for i in range(1, 15)],
    "amount": [100, 50, 200, 150, 75, 300, 125, 80, 90, 60, 200, 45, 110, 95],
})

_D3_df = pd.DataFrame({
    "dept": ["Eng", "Sales", "Eng", "HR", "Sales", "Eng", "HR", "Sales"],
    "employee": ["A", "B", "C", "D", "E", "F", "G", "H"],
    "salary": [120000, 80000, 110000, 70000, 90000, 130000, 75000, 85000],
})

_D4_df = pd.DataFrame({
    "category": ["Electronics", "Clothing", "Electronics", "Food",
                 "Clothing", "Food", "Electronics", "Books", "Books"],
    "product": ["P1", "P2", "P3", "P4", "P5", "P6", "P7", "P8", "P9"],
    "units_sold": [500, 200, 300, 150, 100, 400, 250, 50, 80],
})

_D5_df = pd.DataFrame({
    "venue": ["Arena", "Hall", "Arena", "Park", "Hall", "Arena", "Park", "Hall"],
    "event": ["E1", "E2", "E3", "E4", "E5", "E6", "E7", "E8"],
    "attendees": [500, 200, 300, 100, 250, 400, 150, 300],
    "capacity": [1000, 600, 1000, 200, 600, 1000, 200, 600],
})


# ---------------------------------------------------------------------------
# Compute expected outputs
# ---------------------------------------------------------------------------

# Cluster A
_A1_expected = _A1_df[_A1_df["region"] == "West"].sort_values("revenue", ascending=False)["product"].tolist()
_A2_expected = _A2_df[_A2_df["dept"] == "Eng"].sort_values("tenure", ascending=False)["name"].tolist()
_A3_expected = _A3_df[_A3_df["status"] == "shipped"].sort_values("quantity", ascending=False)["order_id"].tolist()
_A4_expected = _A4_df[_A4_df["grade"] == "A"].sort_values("gpa", ascending=False)["student_id"].tolist()
_A5_expected = _A5_df[_A5_df["stock"] < 10].sort_values("reorder_priority")["item_name"].tolist()

# Cluster B
_B1_expected = _B1_df.assign(email=_B1_df["email"].str.lower())\
    .query("email.str.endswith('@gmail.com')")["email"].tolist()
_B2_expected = _B2_df.assign(name=_B2_df["name"].str.strip())\
    .query("name.str.startswith('A')")["name"].tolist()
_B3_expected = _B3_df.assign(country_code=_B3_df["country_code"].str.upper())\
    .query("country_code == 'US'")["user_id"].tolist()
_B4_expected = _B4_df.assign(product_name=_B4_df["product_name"].str.title())\
    .query("product_name.str.contains('Pro')")["product_id"].tolist()
_B5_expected = _B5_df.assign(phone=_B5_df["phone"].str.replace(r"\D", "", regex=True))\
    .query("phone.str.len() == 10")["phone"].tolist()

# Cluster C
_C1_expected = int(_C1_df.assign(days_since=(pd.Timestamp(_today) - _C1_df["signup_date"]).dt.days)\
    .query("days_since < 30").shape[0])
_C2_expected = int(_C2_df.assign(order_age=(pd.Timestamp(_today) - _C2_df["order_date"]).dt.days)\
    .query("order_age <= 7").shape[0])
_C3_expected = int(_C3_df.assign(tenure_years=(pd.Timestamp(_today) - _C3_df["hire_date"]).dt.days / 365)\
    .query("tenure_years < 2").shape[0])
_C4_expected = int(_C4_df.assign(days_until=(_C4_df["event_date"] - pd.Timestamp(_today)).dt.days)\
    .query("days_until <= 14").shape[0])
_C5_expected = int(_C5_df.assign(age=(pd.Timestamp(_today) - _C5_df["birthdate"]).dt.days / 365)\
    .query("18 <= age <= 25").shape[0])

# Cluster D
_D1_expected = _D1_df.groupby("region")["revenue"].sum()\
    .sort_values(ascending=False).head(3).index.tolist()
_D2_expected = _D2_df.groupby("customer")["order_id"].count()\
    .loc[lambda x: x > 5].index.tolist()
_D3_expected = _D3_df.groupby("dept")["salary"].mean()\
    .sort_values(ascending=False).index[0]
_D4_expected = _D4_df.groupby("category")["units_sold"].sum()\
    .sort_values().head(2).index.tolist()
_D5_expected = _D5_df.assign(over=_D5_df["attendees"] > _D5_df["capacity"])\
    .query("over").groupby("venue")["event"].count().index.tolist()
# D5 recompute: venues where total attendees > total capacity
_D5_agg = _D5_df.groupby("venue").agg({"attendees": "sum", "capacity": "first"}).reset_index()
_D5_expected = _D5_agg[_D5_agg["attendees"] > _D5_agg["capacity"]]["venue"].tolist()


# ---------------------------------------------------------------------------
# TASKS list
# ---------------------------------------------------------------------------

TASKS = [
    # --- Cluster A: filter → sort → extract ---
    {
        "id": "A1",
        "cluster": "A",
        "description": "Given a sales dataframe with columns [product, revenue, region], "
                       "return the product names for the West region sorted by revenue descending.",
        "dataframe": _A1_df,
        "expected_output": _A1_expected,
    },
    {
        "id": "A2",
        "cluster": "A",
        "description": "Given an employees dataframe with columns [name, dept, tenure], "
                       "return the names of Engineering employees sorted by tenure descending.",
        "dataframe": _A2_df,
        "expected_output": _A2_expected,
    },
    {
        "id": "A3",
        "cluster": "A",
        "description": "Given an orders dataframe with columns [order_id, status, quantity], "
                       "return the order IDs for shipped orders sorted by quantity descending.",
        "dataframe": _A3_df,
        "expected_output": _A3_expected,
    },
    {
        "id": "A4",
        "cluster": "A",
        "description": "Given a students dataframe with columns [student_id, grade, gpa], "
                       "return the student IDs of students with grade A sorted by GPA descending.",
        "dataframe": _A4_df,
        "expected_output": _A4_expected,
    },
    {
        "id": "A5",
        "cluster": "A",
        "description": "Given an inventory dataframe with columns [item_name, stock, reorder_priority], "
                       "return item names where stock is below 10, sorted by reorder priority ascending.",
        "dataframe": _A5_df,
        "expected_output": _A5_expected,
    },
    # --- Cluster B: normalize → filter → extract ---
    {
        "id": "B1",
        "cluster": "B",
        "description": "Given a users dataframe with columns [user_id, email], "
                       "lowercase all emails, keep only @gmail.com addresses, return the email list.",
        "dataframe": _B1_df,
        "expected_output": _B1_expected,
    },
    {
        "id": "B2",
        "cluster": "B",
        "description": "Given a contacts dataframe with columns [id, name], "
                       "strip whitespace from names, keep names starting with 'A', return the name list.",
        "dataframe": _B2_df,
        "expected_output": _B2_expected,
    },
    {
        "id": "B3",
        "cluster": "B",
        "description": "Given a users dataframe with columns [user_id, country_code], "
                       "uppercase all country codes, keep only 'US', return the user_id list.",
        "dataframe": _B3_df,
        "expected_output": _B3_expected,
    },
    {
        "id": "B4",
        "cluster": "B",
        "description": "Given a products dataframe with columns [product_id, product_name], "
                       "title-case all product names, keep those containing 'Pro', return the product_id list.",
        "dataframe": _B4_df,
        "expected_output": _B4_expected,
    },
    {
        "id": "B5",
        "cluster": "B",
        "description": "Given a contacts dataframe with columns [contact_id, phone], "
                       "remove all non-digit characters from phone numbers, keep only 10-digit ones, return the phone list.",
        "dataframe": _B5_df,
        "expected_output": _B5_expected,
    },
    # --- Cluster C: date delta → filter → count ---
    {
        "id": "C1",
        "cluster": "C",
        "description": "Given a users dataframe with columns [user_id, signup_date, active], "
                       "compute days since signup (from 2026-03-08), keep users who signed up within the last 30 days, return the count.",
        "dataframe": _C1_df,
        "expected_output": _C1_expected,
    },
    {
        "id": "C2",
        "cluster": "C",
        "description": "Given an orders dataframe with columns [order_id, order_date, amount], "
                       "compute order age in days (from 2026-03-08), keep orders within the last 7 days, return the count.",
        "dataframe": _C2_df,
        "expected_output": _C2_expected,
    },
    {
        "id": "C3",
        "cluster": "C",
        "description": "Given an employees dataframe with columns [emp_id, hire_date, dept], "
                       "compute tenure in years (from 2026-03-08), keep employees with less than 2 years, return the count.",
        "dataframe": _C3_df,
        "expected_output": _C3_expected,
    },
    {
        "id": "C4",
        "cluster": "C",
        "description": "Given an events dataframe with columns [event_id, event_date, venue], "
                       "compute days until event (from 2026-03-08), keep events within the next 14 days, return the count.",
        "dataframe": _C4_df,
        "expected_output": _C4_expected,
    },
    {
        "id": "C5",
        "cluster": "C",
        "description": "Given a users dataframe with columns [user_id, birthdate], "
                       "compute age in years (from 2026-03-08), keep users aged 18 to 25 inclusive, return the count.",
        "dataframe": _C5_df,
        "expected_output": _C5_expected,
    },
    # --- Cluster D: groupby → aggregate → sort → slice ---
    {
        "id": "D1",
        "cluster": "D",
        "description": "Given a sales dataframe with columns [region, revenue, product], "
                       "group by region, sum revenue, return the top 3 regions by total revenue.",
        "dataframe": _D1_df,
        "expected_output": _D1_expected,
    },
    {
        "id": "D2",
        "cluster": "D",
        "description": "Given an orders dataframe with columns [customer, order_id, amount], "
                       "group by customer, count orders, return customers with more than 5 orders.",
        "dataframe": _D2_df,
        "expected_output": _D2_expected,
    },
    {
        "id": "D3",
        "cluster": "D",
        "description": "Given an employees dataframe with columns [dept, employee, salary], "
                       "group by department, compute average salary, return the department with the highest average.",
        "dataframe": _D3_df,
        "expected_output": _D3_expected,
    },
    {
        "id": "D4",
        "cluster": "D",
        "description": "Given a products dataframe with columns [category, product, units_sold], "
                       "group by category, sum units_sold, return the bottom 2 categories by total units.",
        "dataframe": _D4_df,
        "expected_output": _D4_expected,
    },
    {
        "id": "D5",
        "cluster": "D",
        "description": "Given an events dataframe with columns [venue, event, attendees, capacity], "
                       "group by venue summing attendees and taking the first capacity, "
                       "return venues where total attendees exceed capacity.",
        "dataframe": _D5_df,
        "expected_output": _D5_expected,
    },
]