server/dataset_factory.py

# """
# Dataset Factory
# ===============
# Generates deterministic dirty + expected DataFrames for each task.
# Uses numpy default_rng(seed) — every (task_id, seed) pair is identical.

# Task 4 is the novel one: also exposes generate_drift_batch() which the
# environment calls every DRIFT_EVERY steps to inject fresh dirty rows
# mid-episode, simulating a live streaming pipeline under data drift.
# """
# import pandas as pd
# import numpy as np
# from typing import Tuple, Dict


# def make_task(task_id: str, seed: int) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
#     if task_id == "task1":
#         return _task1(seed)
#     elif task_id == "task2":
#         return _task2(seed)
#     elif task_id == "task3":
#         return _task3(seed)
#     elif task_id == "task4_data_drift":
#         return _task4(seed)
#     raise ValueError(f"Unknown task_id: {task_id!r}.")


# # ── Task 1 ────────────────────────────────────────────────────────────────────

# def _task1(seed: int):
#     rng = np.random.default_rng(seed)
#     n = 50
#     ids   = list(range(1, n + 1))
#     names = [f"Customer_{i:03d}" for i in range(n)]
#     ages  = rng.integers(18, 75, size=n).tolist()
#     sals  = np.round(rng.uniform(30_000, 120_000, size=n), 2).tolist()
#     cities = rng.choice(["Mumbai","Delhi","Bangalore","Chennai","Pune"], size=n).tolist()

#     null_age = set(rng.choice(n, size=10, replace=False).tolist())
#     null_sal = set(rng.choice(n, size=8,  replace=False).tolist())
#     markers  = ["", "N/A", "null", "missing", "NaN"]

#     age_d = [str(ages[i]) if i not in null_age else str(rng.choice(markers)) for i in range(n)]
#     sal_d = [sals[i] if i not in null_sal else None for i in range(n)]

#     dirty = pd.DataFrame({"id": ids, "name": names, "age": age_d, "salary": sal_d, "city": cities})

#     age_fill = int(np.median([ages[i] for i in range(n) if i not in null_age]))
#     sal_fill = round(float(np.mean([sals[i] for i in range(n) if i not in null_sal])), 2)

#     expected = pd.DataFrame({
#         "id": ids, "name": names,
#         "age":    pd.array([ages[i] if i not in null_age else age_fill for i in range(n)], dtype="int64"),
#         "salary": pd.array([round(sals[i],2) if i not in null_sal else sal_fill for i in range(n)], dtype="float64"),
#         "city": cities,
#     })
#     return {"main": dirty}, {"main": expected}


# # ── Task 2 ────────────────────────────────────────────────────────────────────

# def _task2(seed: int):
#     rng = np.random.default_rng(seed)
#     nu = 170

#     ids  = list(range(1, nu + 1))
#     cids = rng.integers(1001, 1200, size=nu).tolist()
#     amts = np.round(rng.uniform(10, 5_000, size=nu), 2).tolist()
#     null_a = set(rng.choice(nu, size=12, replace=False).tolist())
#     stats  = rng.choice(["completed","pending","cancelled","refunded"], size=nu).tolist()
#     cats   = rng.choice(["Electronics","Clothing","Food","Books","Sports"], size=nu).tolist()

#     CV = {
#         "USA":["USA","usa","U.S.A","United States","US"],
#         "UK":["UK","uk","U.K.","United Kingdom"],
#         "INDIA":["India","india","INDIA","IN"],
#         "GERMANY":["Germany","germany","DE","GERMANY"],
#         "FRANCE":["France","france","FR","FRANCE"],
#     }
#     ckeys   = list(CV.keys())
#     cc      = rng.choice(ckeys, size=nu).tolist()
#     cd      = [str(rng.choice(CV[c])) for c in cc]

#     dr   = pd.date_range("2023-01-01","2024-12-31", periods=nu)
#     diso = dr.strftime("%Y-%m-%d").tolist()
#     dd   = [pd.Timestamp(d).strftime("%d/%m/%Y") if rng.random()<0.35 else d for d in diso]
#     ad   = [amts[i] if i not in null_a else None for i in range(nu)]

#     base = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cd,
#                          "amount":ad,"order_date":dd,"status":stats,"product_category":cats})
#     dups = base.iloc[rng.choice(nu, size=30, replace=True)].copy()
#     dirty = (pd.concat([base, dups], ignore_index=True)
#              .sample(frac=1, random_state=int(seed)).reset_index(drop=True))

#     af = round(float(np.mean([amts[i] for i in range(nu) if i not in null_a])), 2)
#     expected = pd.DataFrame({
#         "order_id":ids,"customer_id":cids,"country":cc,
#         "amount":pd.array([round(amts[i],2) if i not in null_a else af for i in range(nu)], dtype="float64"),
#         "order_date":pd.to_datetime(diso),"status":stats,"product_category":cats,
#     })
#     return {"main": dirty}, {"main": expected}


# # ── Task 3 ────────────────────────────────────────────────────────────────────

# def _task3(seed: int):
#     rng = np.random.default_rng(seed)
#     nc, no = 100, 300

#     cids  = list(range(1001, 1001+nc))
#     cname = [f"Customer_{i:03d}" for i in range(nc)]
#     ctry  = rng.choice(["USA","UK","India","Germany"], size=nc).tolist()
#     ages  = rng.integers(18, 70, size=nc).tolist()
#     nai   = set(rng.choice(nc, size=8, replace=False).tolist())
#     ages_d = [str(ages[i]) if i not in nai else "N/A" for i in range(nc)]

#     cust_dirty = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,"age":ages_d})
#     af = int(np.median([ages[i] for i in range(nc) if i not in nai]))
#     cust_clean = pd.DataFrame({
#         "customer_id":cids,"name":cname,"country":ctry,
#         "age":pd.array([ages[i] if i not in nai else af for i in range(nc)], dtype="int64"),
#     })

#     oids = list(range(1, no+1))
#     ocid = rng.choice(cids, size=no).tolist()
#     amts = np.round(rng.uniform(10, 2_000, size=no), 2)
#     for idx in rng.choice(no, size=20, replace=False):
#         amts[idx] = float(rng.choice([0.01, -5.0, 50_000.0, 99_999.0]))
#     dates = pd.date_range("2023-01-01","2024-12-31", periods=no).strftime("%Y-%m-%d").tolist()

#     orders_dirty = pd.DataFrame({"order_id":oids,"customer_id":ocid,
#                                  "amount":amts.tolist(),"order_date":dates})
#     merged = pd.merge(orders_dirty, cust_clean, on="customer_id", how="inner")
#     Q1, Q3 = merged["amount"].quantile(0.25), merged["amount"].quantile(0.75)
#     IQR = Q3 - Q1
#     mc = merged[(merged["amount"]>=Q1-1.5*IQR) & (merged["amount"]<=Q3+1.5*IQR)].copy().reset_index(drop=True)
#     mc["order_year"] = pd.to_datetime(mc["order_date"]).dt.year

#     return ({"orders": orders_dirty, "customers": cust_dirty}, {"main": mc})


# # ── Task 4: Data Drift (Expert) ───────────────────────────────────────────────

# def _task4(seed: int):
#     """
#     Live streaming transactions — 120 initial dirty rows.
#     Env injects fresh dirty rows every DRIFT_EVERY=5 steps via generate_drift_batch().
#     Agent must keep cleaning as new dirty data continuously arrives.

#     Columns: txn_id, customer_id, amount, category, region, event_ts
#     Dirty issues: nulls, wrong dtypes (amount as str), outliers, mixed timestamp formats.
#     """
#     rng = np.random.default_rng(seed)
#     n = 120

#     txn_ids  = [f"TXN_INIT_{i:04d}" for i in range(n)]
#     cids     = rng.integers(1, 501, size=n).tolist()
#     cats_c   = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n).tolist()
#     regs_c   = rng.choice(["North","South","East","West","Central"], size=n).tolist()
#     amts_t   = np.round(rng.uniform(10, 3000, size=n), 2).tolist()

#     amts_d = []
#     for i in range(n):
#         r = rng.random()
#         if r < 0.15:   amts_d.append(None)
#         elif r < 0.22: amts_d.append(str(round(amts_t[i], 2)))
#         elif r < 0.27: amts_d.append(float(-rng.uniform(100, 5000)))
#         elif r < 0.31: amts_d.append(float(rng.uniform(80000, 250000)))
#         else:          amts_d.append(amts_t[i])

#     def _ts(rng):
#         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
#                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
#         r = rng.random()
#         if r < 0.20: return base.split(" ")[0].replace("-", "/")
#         if r < 0.30:
#             p = base.split("-"); return f"{p[2][:2]}/{p[1]}/{p[0]}"
#         return base

#     ts_d   = [_ts(rng) for _ in range(n)]
#     cats_d = [None if rng.random()<0.15 else cats_c[i] for i in range(n)]
#     regs_d = [None if rng.random()<0.10 else regs_c[i] for i in range(n)]

#     dirty = pd.DataFrame({
#         "txn_id":txn_ids, "customer_id":cids,
#         "amount":amts_d, "category":cats_d, "region":regs_d, "event_ts":ts_d,
#     })

#     # Expected: cleaned initial batch (outliers dropped, nulls filled, ts parsed)
#     good_amts = [x for x in amts_t if 0 < x <= 10000]
#     amt_fill  = round(float(np.mean(good_amts)), 2)
#     amts_e = []
#     for a in amts_d:
#         if a is None:                               amts_e.append(amt_fill)
#         elif isinstance(a, str):                    amts_e.append(float(a))
#         elif isinstance(a, float) and (a<0 or a>10000): amts_e.append(None)
#         else:                                       amts_e.append(round(a, 2))

#     exp_df = pd.DataFrame({
#         "txn_id":txn_ids, "customer_id":cids,
#         "amount":pd.to_numeric(amts_e, errors="coerce"),
#         "category":cats_c, "region":regs_c,
#         "event_ts":pd.to_datetime(ts_d, errors="coerce"),
#     }).dropna(subset=["amount"]).reset_index(drop=True)

#     return {"stream": dirty}, {"stream": exp_df}


# # ── Drift Batch Generator ─────────────────────────────────────────────────────

# def generate_drift_batch(seed: int, batch_num: int, n_rows: int = 7) -> pd.DataFrame:
#     """
#     Generate a fresh batch of dirty rows injected mid-episode into task4.
#     Called by DataCleanEnvironment.step() every DRIFT_EVERY steps.

#     Fully deterministic: (seed, batch_num) always → same batch.
#     Each batch introduces different dirty patterns so the agent faces novel problems.
#     """
#     rng = np.random.default_rng(seed * 1000 + batch_num)

#     txn_ids  = [f"TXN_DRIFT_{batch_num:03d}_{i:02d}" for i in range(n_rows)]
#     cids     = rng.integers(1, 501, size=n_rows).tolist()
#     cats     = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n_rows).tolist()
#     regs     = rng.choice(["North","South","East","West","Central"], size=n_rows).tolist()

#     amts = []
#     for _ in range(n_rows):
#         r    = rng.random()
#         base = round(float(rng.uniform(10, 3000)), 2)
#         if r < 0.20:   amts.append(None)
#         elif r < 0.30: amts.append(str(base))
#         elif r < 0.38: amts.append(float(-rng.uniform(100, 5000)))
#         elif r < 0.44: amts.append(float(rng.uniform(80000, 250000)))
#         else:          amts.append(base)

#     ts = []
#     for _ in range(n_rows):
#         base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
#                 f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
#         r = rng.random()
#         if r < 0.20:   ts.append(base.split(" ")[0].replace("-", "/"))
#         elif r < 0.35: p = base.split("-"); ts.append(f"{p[2][:2]}/{p[1]}/{p[0]}")
#         else:          ts.append(base)

#     for i in range(n_rows):
#         if rng.random() < 0.18: cats[i] = None
#         if rng.random() < 0.12: regs[i] = None

#     return pd.DataFrame({
#         "txn_id":txn_ids, "customer_id":cids,
#         "amount":amts, "category":cats, "region":regs, "event_ts":ts,
#     })

import pandas as pd
import numpy as np
from typing import Tuple, Dict

def make_task(task_id: str, seed: int) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
    if task_id == "task1": return _task1(seed)
    elif task_id == "task2": return _task2(seed)
    elif task_id == "task3": return _task3(seed)
    elif task_id == "task4_data_drift": return _task4(seed)
    raise ValueError(f"Unknown task_id: {task_id!r}.")

def _task1(seed):
    rng = np.random.default_rng(seed)
    n = 50
    ids = list(range(1, n+1))
    names = [f"Customer_{i:03d}" for i in range(n)]
    ages = rng.integers(18, 75, size=n).tolist()
    sals = np.round(rng.uniform(30000, 120000, size=n), 2).tolist()
    cities = rng.choice(["Mumbai","Delhi","Bangalore","Chennai","Pune"], size=n).tolist()
    null_age = set(rng.choice(n, size=10, replace=False).tolist())
    null_sal = set(rng.choice(n, size=8, replace=False).tolist())
    markers = ["", "N/A", "null", "missing", "NaN"]
    age_d = [str(ages[i]) if i not in null_age else str(rng.choice(markers)) for i in range(n)]
    sal_d = [sals[i] if i not in null_sal else None for i in range(n)]
    dirty = pd.DataFrame({"id":ids,"name":names,"age":age_d,"salary":sal_d,"city":cities})
    age_fill = int(np.median([ages[i] for i in range(n) if i not in null_age]))
    sal_fill = round(float(np.mean([sals[i] for i in range(n) if i not in null_sal])), 2)
    expected = pd.DataFrame({
        "id":ids,"name":names,
        "age": pd.array([ages[i] if i not in null_age else age_fill for i in range(n)], dtype="int64"),
        "salary": pd.array([round(sals[i],2) if i not in null_sal else sal_fill for i in range(n)], dtype="float64"),
        "city":cities,
    })
    return {"main": dirty}, {"main": expected}

def _task2(seed):
    rng = np.random.default_rng(seed)
    nu = 170
    ids = list(range(1, nu+1))
    cids = rng.integers(1001, 1200, size=nu).tolist()
    amts = np.round(rng.uniform(10, 5000, size=nu), 2).tolist()
    null_a = set(rng.choice(nu, size=12, replace=False).tolist())
    stats = rng.choice(["completed","pending","cancelled","refunded"], size=nu).tolist()
    cats = rng.choice(["Electronics","Clothing","Food","Books","Sports"], size=nu).tolist()
    CV = {"USA":["USA","usa","U.S.A","United States","US"],"UK":["UK","uk","U.K.","United Kingdom"],
          "INDIA":["India","india","INDIA","IN"],"GERMANY":["Germany","germany","DE","GERMANY"],
          "FRANCE":["France","france","FR","FRANCE"]}
    ckeys = list(CV.keys())
    cc = rng.choice(ckeys, size=nu).tolist()
    cd = [str(rng.choice(CV[c])) for c in cc]
    dr = pd.date_range("2023-01-01","2024-12-31", periods=nu)
    diso = dr.strftime("%Y-%m-%d").tolist()
    dd = [pd.Timestamp(d).strftime("%d/%m/%Y") if rng.random()<0.35 else d for d in diso]
    ad = [amts[i] if i not in null_a else None for i in range(nu)]
    base = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cd,"amount":ad,"order_date":dd,"status":stats,"product_category":cats})
    dups = base.iloc[rng.choice(nu, size=30, replace=True)].copy()
    dirty = (pd.concat([base, dups], ignore_index=True).sample(frac=1, random_state=int(seed)).reset_index(drop=True))
    af = round(float(np.mean([amts[i] for i in range(nu) if i not in null_a])), 2)
    expected = pd.DataFrame({"order_id":ids,"customer_id":cids,"country":cc,
        "amount":pd.array([round(amts[i],2) if i not in null_a else af for i in range(nu)], dtype="float64"),
        "order_date":pd.to_datetime(diso),"status":stats,"product_category":cats})
    return {"main": dirty}, {"main": expected}

def _task3(seed):
    rng = np.random.default_rng(seed)
    nc, no = 100, 300
    cids = list(range(1001, 1001+nc))
    cname = [f"Customer_{i:03d}" for i in range(nc)]
    ctry = rng.choice(["USA","UK","India","Germany"], size=nc).tolist()
    ages = rng.integers(18, 70, size=nc).tolist()
    nai = set(rng.choice(nc, size=8, replace=False).tolist())
    ages_d = [str(ages[i]) if i not in nai else "N/A" for i in range(nc)]
    cust_dirty = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,"age":ages_d})
    af = int(np.median([ages[i] for i in range(nc) if i not in nai]))
    cust_clean = pd.DataFrame({"customer_id":cids,"name":cname,"country":ctry,
        "age":pd.array([ages[i] if i not in nai else af for i in range(nc)], dtype="int64")})
    oids = list(range(1, no+1))
    ocid = rng.choice(cids, size=no).tolist()
    amts = np.round(rng.uniform(10, 2000, size=no), 2)
    for idx in rng.choice(no, size=20, replace=False):
        amts[idx] = float(rng.choice([0.01, -5.0, 50000.0, 99999.0]))
    dates = pd.date_range("2023-01-01","2024-12-31", periods=no).strftime("%Y-%m-%d").tolist()
    orders_dirty = pd.DataFrame({"order_id":oids,"customer_id":ocid,"amount":amts.tolist(),"order_date":dates})
    merged = pd.merge(orders_dirty, cust_clean, on="customer_id", how="inner")
    Q1, Q3 = merged["amount"].quantile(0.25), merged["amount"].quantile(0.75)
    IQR = Q3 - Q1
    mc = merged[(merged["amount"]>=Q1-1.5*IQR) & (merged["amount"]<=Q3+1.5*IQR)].copy().reset_index(drop=True)
    mc["order_year"] = pd.to_datetime(mc["order_date"]).dt.year
    return ({"orders": orders_dirty, "customers": cust_dirty}, {"main": mc})

def _task4(seed):
    rng = np.random.default_rng(seed)
    n = 120
    txn_ids = [f"TXN_INIT_{i:04d}" for i in range(n)]
    cids = rng.integers(1, 501, size=n).tolist()
    cats_c = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n).tolist()
    regs_c = rng.choice(["North","South","East","West","Central"], size=n).tolist()
    amts_t = np.round(rng.uniform(10, 3000, size=n), 2).tolist()
    amts_d = []
    for i in range(n):
        r = rng.random()
        if r < 0.15: amts_d.append(None)
        elif r < 0.22: amts_d.append(str(round(amts_t[i], 2)))
        elif r < 0.27: amts_d.append(float(-rng.uniform(100, 5000)))
        elif r < 0.31: amts_d.append(float(rng.uniform(80000, 250000)))
        else: amts_d.append(amts_t[i])
    def _ts(rng):
        base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
                f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
        r = rng.random()
        if r < 0.20: return base.split(" ")[0].replace("-", "/")
        if r < 0.30:
            p = base.split("-"); return f"{p[2][:2]}/{p[1]}/{p[0]}"
        return base
    ts_d = [_ts(rng) for _ in range(n)]
    cats_d = [None if rng.random()<0.15 else cats_c[i] for i in range(n)]
    regs_d = [None if rng.random()<0.10 else regs_c[i] for i in range(n)]
    dirty = pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,"amount":amts_d,"category":cats_d,"region":regs_d,"event_ts":ts_d})
    good_amts = [x for x in amts_t if 0 < x <= 10000]
    amt_fill = round(float(np.mean(good_amts)), 2)
    amts_e = []
    for a in amts_d:
        if a is None: amts_e.append(amt_fill)
        elif isinstance(a, str): amts_e.append(float(a))
        elif isinstance(a, float) and (a<0 or a>10000): amts_e.append(None)
        else: amts_e.append(round(a, 2))
    exp_df = pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,
        "amount":pd.to_numeric(amts_e, errors="coerce"),"category":cats_c,
        "region":regs_c,"event_ts":pd.to_datetime(ts_d, errors="coerce")}).dropna(subset=["amount"]).reset_index(drop=True)
    return {"stream": dirty}, {"stream": exp_df}

def generate_drift_batch(seed: int, batch_num: int, n_rows: int = 7) -> pd.DataFrame:
    rng = np.random.default_rng(seed * 1000 + batch_num)
    txn_ids = [f"TXN_DRIFT_{batch_num:03d}_{i:02d}" for i in range(n_rows)]
    cids = rng.integers(1, 501, size=n_rows).tolist()
    cats = rng.choice(["Electronics","Clothing","Food","Books","Sports","Toys"], size=n_rows).tolist()
    regs = rng.choice(["North","South","East","West","Central"], size=n_rows).tolist()
    amts = []
    for _ in range(n_rows):
        r = rng.random(); base = round(float(rng.uniform(10, 3000)), 2)
        if r < 0.20: amts.append(None)
        elif r < 0.30: amts.append(str(base))
        elif r < 0.38: amts.append(float(-rng.uniform(100, 5000)))
        elif r < 0.44: amts.append(float(rng.uniform(80000, 250000)))
        else: amts.append(base)
    ts = []
    for _ in range(n_rows):
        base = (f"2024-{rng.integers(1,13):02d}-{rng.integers(1,29):02d} "
                f"{rng.integers(0,24):02d}:{rng.integers(0,60):02d}:00")
        r = rng.random()
        if r < 0.20: ts.append(base.split(" ")[0].replace("-", "/"))
        elif r < 0.35:
            p = base.split("-"); ts.append(f"{p[2][:2]}/{p[1]}/{p[0]}")
        else: ts.append(base)
    for i in range(n_rows):
        if rng.random() < 0.18: cats[i] = None
        if rng.random() < 0.12: regs[i] = None
    return pd.DataFrame({"txn_id":txn_ids,"customer_id":cids,"amount":amts,"category":cats,"region":regs,"event_ts":ts})