Data

"""Deterministic fictional depot jobs. No real people or company records."""
from pathlib import Path
import hashlib
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
SEED = 20261014
NUMERIC = ['distance_km', 'queued_jobs', 'item_count', 'weekday', 'promised_hours']
CATEGORICAL = ['priority', 'zone', 'depot_id']
FEATURES = NUMERIC + CATEGORICAL
FORBIDDEN = ['duration_hours', 'late', 'label_available_at', 'final_delay_reason', 'close_cost']
DATE_COLS = ['intake_at', 'label_available_at']

def generate(seed=SEED):
    """1500 rows, 6 jobs/day; targets available within 24 h of intake."""
    rng = np.random.default_rng(seed)
    n = 1500
    day = np.repeat(np.arange(250), 6)
    hour = np.tile(np.arange(8, 20, 2), 250)
    intake = pd.Timestamp('2025-01-01', tz='UTC') + pd.to_timedelta(day, unit='D') + pd.to_timedelta(hour, unit='h')
    depot_n = rng.integers(0, 12, n)
    distance = np.round(rng.uniform(2, 80, n), 2)
    queue = rng.poisson(6 + 2 * (day >= 190), n)
    items = rng.integers(1, 11, n)
    priority = rng.choice(['standard', 'express'], n, p=[0.75, 0.25])
    zone = rng.choice(['central', 'east', 'west'], n, p=[0.45, 0.3, 0.25])
    promised = np.where(priority == 'express', 6., 10.)
    depot_effect = np.array([-0.7, 0.4, 0.8, -0.2, 0.1, -0.5, 0.6, -0.4, 0.3, 0.5, -0.1, 0.2])[depot_n]
    # Queue congestion and a gradual mechanism change are deliberately synthetic.
    mechanism_shift = np.maximum(day - 180, 0) * 0.012
    duration = (1.5 + 0.06 * distance + 0.30 * queue + 0.17 * items
                + 0.10 * np.maximum(queue - 7, 0) ** 2
                + 0.8 * (zone == 'west') - 0.7 * (priority == 'express')
                + depot_effect + mechanism_shift + rng.normal(0, 1.0, n))
    duration = np.round(np.clip(duration, 0.5, 23.5), 3)
    frame = pd.DataFrame({
        'job_id': [f'J{i:05d}' for i in range(n)], 'day': day, 'intake_at': intake,
        'distance_km': distance, 'queued_jobs': queue, 'item_count': items,
        'priority': priority, 'zone': zone,
        'depot_id': [f'D{i:02d}' for i in depot_n],
        'weekday': intake.dayofweek, 'promised_hours': promised,
        'duration_hours': duration, 'late': (duration > promised).astype(int),
        'label_available_at': intake + pd.to_timedelta(duration, unit='h'),
        'final_delay_reason': np.where(duration > promised, 'late_confirmed', 'on_time_confirmed'),
        'close_cost': np.round(12 + 3 * duration + rng.normal(0, .3, n), 2)
    })
    # Missing inputs are introduced after the outcome mechanism has used true values.
    frame.loc[rng.random(n) < .05, 'distance_km'] = np.nan
    frame.loc[rng.random(n) < .03, 'queued_jobs'] = np.nan
    return frame

def parts(frame):
    return {
        'train': frame[frame.day < 149].copy(),
        'validation': frame[(frame.day >= 150) & (frame.day < 199)].copy(),
        'test': frame[frame.day >= 200].copy(),
        'embargo': frame[frame.day.isin([149, 199])].copy(),
    }

def write_data():
    folder = ROOT / 'data'
    folder.mkdir(exist_ok=True)
    for name, frame in parts(generate()).items():
        frame.to_csv(folder / f'{name}.csv', index=False, float_format='%.6f')
    return {k: len(v) for k, v in parts(generate()).items()}

def load_part(name):
    if name not in {'train', 'validation', 'test', 'embargo'}:
        raise ValueError('Unknown data partition')
    frame = pd.read_csv(ROOT / 'data' / f'{name}.csv')
    for column in DATE_COLS:
        frame[column] = pd.to_datetime(frame[column], format='mixed', utc=True, errors='raise')
    return frame

def validation_roles(validation):
    return (validation[validation.day < 174].copy(),
            validation[validation.day >= 175].copy())

def file_hash(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()

def assert_available(before, after):
    if not before.label_available_at.max() < after.intake_at.min():
        raise ValueError('A training/calibration label was unavailable at prediction time')

def audit(frame):
    assert frame.job_id.is_unique
    assert frame.intake_at.is_monotonic_increasing
    assert frame.late.isin([0, 1]).all()
    assert (frame.late == (frame.duration_hours > frame.promised_hours)).all()
    assert ((frame.label_available_at - frame.intake_at).dt.total_seconds() > 0).all()
    assert ((frame.label_available_at - frame.intake_at).dt.total_seconds() < 86400).all()
    assert frame[NUMERIC].notna().any().all()
    assert not set(FEATURES) & set(FORBIDDEN)
    return {'rows': len(frame), 'late_rate': float(frame.late.mean()),
            'missing_inputs': {c: int(frame[c].isna().sum()) for c in FEATURES}}

if __name__ == '__main__':
    print(write_data())