src/data_generation.py

"""
data_generation.py
-------------------
Generates the synthetic datasets used to train/evaluate the two ML models
that power the Smart Warehouse AI Assistant:

1. Intent classifier   -> routes free-text queries into warehouse-ops intents
2. Anomaly detector     -> flags abnormal conveyor/AGV sensor readings

All data is synthetically generated with templates + randomised slots so the
project is fully self-contained and reproducible (no external datasets or
scraping required). A fixed random seed keeps results reproducible.
"""

import random
import numpy as np
import pandas as pd

RANDOM_SEED = 42


# --------------------------------------------------------------------------
# 1. INTENT CLASSIFICATION DATA
# --------------------------------------------------------------------------

INTENT_TEMPLATES = {
    "inventory_check": [
        "How many units of {sku} are in {zone}?",
        "What is the current stock level for {sku}?",
        "Check inventory count for {sku} in {zone}",
        "Do we have enough {sku} to fulfill 200 units?",
        "Show me the on-hand quantity of {sku}",
        "Is {sku} in stock at {zone}?",
        "Give me stock levels across all zones for {sku}",
        "How much inventory is left for {sku}?",
    ],
    "order_status": [
        "What's the status of order {order_id}?",
        "Has order {order_id} shipped yet?",
        "Track order {order_id} for me",
        "Is order {order_id} delayed?",
        "When will order {order_id} be delivered?",
        "Show the fulfillment status of {order_id}",
        "Why hasn't order {order_id} left the dock yet?",
    ],
    "equipment_maintenance": [
        "The conveyor belt in {zone} is making noise",
        "Crane {equip_id} reported a fault code",
        "Schedule maintenance for {equip_id}",
        "{equip_id} motor temperature seems high",
        "The sorter in {zone} keeps jamming",
        "Report vibration issue on {equip_id}",
        "Belt {equip_id} stopped unexpectedly, please check",
        "Log a breakdown for {equip_id} in {zone}",
    ],
    "agv_navigation": [
        "Route {equip_id} to picking station {station}",
        "Send AGV {equip_id} to {zone}",
        "Why is {equip_id} stuck near {zone}?",
        "Reassign {equip_id} to charging station",
        "What is the current location of {equip_id}?",
        "Redirect {equip_id} around the blocked aisle in {zone}",
    ],
    "picking_optimization": [
        "What's the fastest picking route for order {order_id}?",
        "Optimize the pick path for {zone}",
        "Should we batch pick these orders together?",
        "Suggest a wave picking plan for {zone}",
        "How can we reduce travel time for pickers in {zone}?",
        "Recommend a picking strategy for high-velocity SKUs",
    ],
    "safety_incident": [
        "A forklift near-miss was reported in {zone}",
        "Log a safety incident involving {equip_id}",
        "There was a near collision between {equip_id} and a pedestrian in {zone}",
        "File an incident report for {zone}",
        "A worker slipped near {equip_id}, please log it",
        "Report unsafe pallet stacking in {zone}",
    ],
    "system_status": [
        "Is {equip_id} operational?",
        "What is the uptime for {equip_id} today?",
        "Check system health for {zone}",
        "Are all cranes online in {zone}?",
        "Give me the current status of the WMS integration",
        "Is the sorter in {zone} running normally?",
    ],
    "general_faq": [
        "What is a WMS?",
        "Explain how an AS/RS works",
        "What's the difference between AGV and AMR?",
        "What is cycle counting?",
        "How does goods-to-person picking work?",
        "What KPIs matter most in warehouse automation?",
        "What is predictive maintenance?",
        "How do sortation systems decide where to route a parcel?",
    ],
}

ZONES = ["Zone A", "Zone B", "Zone C", "Zone D", "the mezzanine", "the receiving dock"]
EQUIP_IDS = ["AGV-07", "AGV-12", "Crane-03", "Sorter-02", "Conveyor-14", "AMR-21", "Crane-05"]
STATIONS = ["3", "5", "7", "12"]


def _rand_sku():
    return f"SKU-{random.randint(1000, 9999)}"


def _rand_order():
    return f"#{random.randint(10000, 99999)}"


def generate_intent_dataset(n_per_intent: int = 45, seed: int = RANDOM_SEED) -> pd.DataFrame:
    """Generate a labelled (text, intent) dataset by sampling + slot-filling templates."""
    rng = random.Random(seed)
    rows = []
    for intent, templates in INTENT_TEMPLATES.items():
        for _ in range(n_per_intent):
            template = rng.choice(templates)
            text = template.format(
                sku=_rand_sku(),
                order_id=_rand_order(),
                zone=rng.choice(ZONES),
                equip_id=rng.choice(EQUIP_IDS),
                station=rng.choice(STATIONS),
            )
            rows.append({"text": text, "intent": intent})
    df = pd.DataFrame(rows)
    df = df.sample(frac=1.0, random_state=seed).reset_index(drop=True)
    return df


# --------------------------------------------------------------------------
# 2. INVENTORY / ORDERS DATA (used by the Inventory & Task Query tab)
# --------------------------------------------------------------------------

def generate_inventory_db(n_skus: int = 60, seed: int = RANDOM_SEED) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    categories = ["Electronics", "Apparel", "Automotive Parts", "Food & Beverage", "Household"]
    zones = ["Zone A", "Zone B", "Zone C", "Zone D"]
    rows = []
    for i in range(n_skus):
        sku = f"SKU-{1000 + i}"
        rows.append({
            "sku": sku,
            "description": f"{rng.choice(categories)} item {1000 + i}",
            "category": rng.choice(categories),
            "zone": rng.choice(zones),
            "on_hand_units": int(rng.integers(0, 2000)),
            "reorder_point": int(rng.integers(50, 300)),
            "unit_cost_jpy": int(rng.integers(200, 15000)),
        })
    return pd.DataFrame(rows)


def generate_orders_db(n_orders: int = 80, seed: int = RANDOM_SEED) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    statuses = ["Received", "Picking", "Packed", "Shipped", "Delayed"]
    weights = [0.15, 0.25, 0.2, 0.3, 0.1]
    rows = []
    for i in range(n_orders):
        order_id = f"#{10000 + i}"
        rows.append({
            "order_id": order_id,
            "status": rng.choice(statuses, p=weights),
            "num_lines": int(rng.integers(1, 25)),
            "priority": rng.choice(["Standard", "Express", "Same-Day"], p=[0.6, 0.3, 0.1]),
            "zone": rng.choice(["Zone A", "Zone B", "Zone C", "Zone D"]),
        })
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------
# 3. SENSOR DATA FOR ANOMALY DETECTION (predictive maintenance)
# --------------------------------------------------------------------------

def generate_sensor_dataset(n_normal: int = 900, n_anomaly: int = 100, seed: int = RANDOM_SEED) -> pd.DataFrame:
    """
    Synthetic conveyor/crane motor sensor readings.
    Features: motor_temp_c, vibration_mm_s, current_amps, belt_speed_mps
    Label: 1 = anomaly (bearing wear / misalignment / overload pattern), 0 = normal
    """
    rng = np.random.default_rng(seed)

    normal = pd.DataFrame({
        "motor_temp_c": rng.normal(55, 4, n_normal).clip(35, 75),
        "vibration_mm_s": rng.normal(2.2, 0.5, n_normal).clip(0.2, 5),
        "current_amps": rng.normal(12, 1.5, n_normal).clip(5, 20),
        "belt_speed_mps": rng.normal(1.5, 0.15, n_normal).clip(0.8, 2.2),
        "label": 0,
    })

    # Anomalies: elevated temp + vibration + current, reduced/erratic belt speed
    anomaly = pd.DataFrame({
        "motor_temp_c": rng.normal(78, 6, n_anomaly).clip(65, 100),
        "vibration_mm_s": rng.normal(5.5, 1.2, n_anomaly).clip(3.5, 10),
        "current_amps": rng.normal(19, 2.5, n_anomaly).clip(14, 28),
        "belt_speed_mps": rng.normal(0.9, 0.3, n_anomaly).clip(0.1, 1.6),
        "label": 1,
    })

    df = pd.concat([normal, anomaly], ignore_index=True)
    df = df.sample(frac=1.0, random_state=seed).reset_index(drop=True)
    return df


# --------------------------------------------------------------------------
# 4. RETRIEVAL EVALUATION SET (query -> expected KB doc id)
# --------------------------------------------------------------------------

RETRIEVAL_EVAL_SET = [
    ("How does an AS/RS crane retrieve a pallet?", "asrs_overview"),
    ("What's the difference between an AGV and an AMR?", "agv_amr_overview"),
    ("What does a WMS integrate with?", "wms_overview"),
    ("Why would a sorter jam?", "conveyor_sorting"),
    ("What is batch picking?", "picking_strategies"),
    ("How can we predict a motor failure before it happens?", "predictive_maintenance"),
    ("What should I do after a near-miss with a forklift?", "safety_protocol"),
    ("How do we keep inventory counts accurate?", "inventory_accuracy"),
    ("What KPIs should a warehouse manager track?", "kpi_overview"),
    ("How can automated warehouses save energy?", "energy_efficiency"),
]