src/data_generator.py

"""
data_generator.py — Training Data Generation for DAHS_2

NEW in DAHS_2: Snapshot-fork algorithm
  Instead of running full simulations with each heuristic,
  this generator takes snapshots every 10 minutes, forks 6 short
  simulations (20 min each), and labels which heuristic wins per-window.
  Result: ~60 rows per scenario instead of 1, with situation-level labels.

Also generates:
  - priority_dataset.csv (same as DAHS_1)
"""

from __future__ import annotations

import logging
import multiprocessing as mp
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from tqdm import tqdm

logger = logging.getLogger(__name__)

DATA_DIR = Path(__file__).parent.parent / "data" / "raw"

HEURISTIC_NAMES = [
    "fifo",
    "priority_edd",
    "critical_ratio",
    "atc",
    "wspt",
    "slack",
]

SNAPSHOT_INTERVAL = 15.0   # minutes between snapshots (matches BatchwiseSelector.EVAL_INTERVAL)
FORK_WINDOW = 60.0         # minutes per fork evaluation (covers express SLA window of 60 min)


# ---------------------------------------------------------------------------
# 7-region scenario diversity (ported from DAHS_1)
# ---------------------------------------------------------------------------

def _make_diverse_scenario_configs(n_scenarios: int, rng: np.random.Generator) -> List[Dict[str, Any]]:
    """Generate diverse simulator parameter configs to avoid class imbalance."""
    configs: List[Dict[str, Any]] = []

    regions = [
        # FIFO-friendly: low load, uniform jobs, loose deadlines
        {"arrival": (1.0, 2.0), "bkdown": (0.0, 0.001), "due": (1.8, 3.0),
         "batch": (5, 15), "lunch": (1.0, 1.1), "pscale": (0.8, 1.2),
         "mix": "uniform"},
        # Priority-EDD: high express, tight deadlines
        {"arrival": (2.0, 3.5), "bkdown": (0.0, 0.005), "due": (0.4, 0.8),
         "batch": (15, 40), "lunch": (1.0, 1.3), "pscale": (0.8, 1.2),
         "mix": "express_heavy"},
        # Critical-Ratio: high breakdowns, heterogeneous pressure
        {"arrival": (2.0, 3.0), "bkdown": (0.008, 0.020), "due": (0.6, 1.2),
         "batch": (20, 50), "lunch": (1.2, 1.6), "pscale": (1.0, 1.5),
         "mix": "diverse"},
        # ATC: heavy load + surge, weighted tardiness matters
        {"arrival": (3.0, 5.0), "bkdown": (0.001, 0.008), "due": (0.7, 1.1),
         "batch": (30, 80), "lunch": (1.2, 1.5), "pscale": (0.9, 1.3),
         "mix": "diverse"},
        # WSPT: many short jobs, steady flow
        {"arrival": (2.5, 4.0), "bkdown": (0.0, 0.003), "due": (1.0, 1.8),
         "batch": (10, 30), "lunch": (1.0, 1.2), "pscale": (0.5, 0.9),
         "mix": "short_heavy"},
        # Slack: tight deadlines, recovery-mode
        {"arrival": (2.5, 3.5), "bkdown": (0.003, 0.012), "due": (0.2, 0.5),
         "batch": (20, 50), "lunch": (1.3, 1.8), "pscale": (1.0, 1.4),
         "mix": "diverse"},
        # Default / general
        {"arrival": (1.5, 4.0), "bkdown": (0.0, 0.015), "due": (0.5, 2.0),
         "batch": (10, 60), "lunch": (1.0, 1.5), "pscale": (0.7, 1.3),
         "mix": "random"},
    ]

    mix_templates = {
        "uniform": {"A": 0.0, "B": 0.0, "C": 1.0, "D": 0.0, "E": 0.0},
        "express_heavy": {"A": 0.20, "B": 0.10, "C": 0.10, "D": 0.10, "E": 0.50},
        "short_heavy": {"A": 0.35, "B": 0.10, "C": 0.10, "D": 0.05, "E": 0.40},
        "diverse": {"A": 0.25, "B": 0.25, "C": 0.20, "D": 0.15, "E": 0.15},
    }

    per_region = n_scenarios // len(regions)
    remainder = n_scenarios - per_region * len(regions)

    seed_counter = 0
    for ri, region in enumerate(regions):
        count = per_region + (1 if ri < remainder else 0)
        for _ in range(count):
            ar  = rng.uniform(*region["arrival"])
            bk  = rng.uniform(*region["bkdown"])
            dd  = rng.uniform(*region["due"])
            bat = int(rng.uniform(*region["batch"]))
            lp  = rng.uniform(*region["lunch"])
            ps  = rng.uniform(*region["pscale"])

            if region["mix"] == "random":
                freqs_raw = rng.dirichlet([1, 1, 1, 1, 1])
                jt_freq = {k: float(v) for k, v in zip("ABCDE", freqs_raw)}
            elif region["mix"] in mix_templates:
                base = mix_templates[region["mix"]].copy()
                noise = rng.uniform(-0.05, 0.05, 5)
                vals = np.array([base[k] for k in "ABCDE"]) + noise
                vals = np.clip(vals, 0.01, None)
                vals /= vals.sum()
                jt_freq = {k: float(v) for k, v in zip("ABCDE", vals)}
            else:
                jt_freq = {}

            configs.append({
                "seed": seed_counter,
                "base_arrival_rate": round(ar, 2),
                "breakdown_prob": round(bk, 4),
                "batch_arrival_size": bat,
                "lunch_penalty_factor": round(lp, 2),
                "job_type_frequencies": jt_freq,
                "due_date_tightness": round(dd, 2),
                "processing_time_scale": round(ps, 2),
            })
            seed_counter += 1

    return configs


# ---------------------------------------------------------------------------
# NEW: Snapshot-fork worker (top-level for multiprocessing)
# ---------------------------------------------------------------------------

def _run_snapshot_scenario(args: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Worker: run one full scenario with snapshot-fork labeling.

    Algorithm:
    1. Run base sim (FIFO) to each 10-minute snapshot
    2. At each snapshot, save state and fork 6 heuristics 20 min each
    3. Label the snapshot with the best-performing heuristic
    Returns ~60 rows per scenario.
    """
    config = args
    from src.heuristics import (
        fifo_dispatch, priority_edd_dispatch, critical_ratio_dispatch,
        atc_dispatch, wspt_dispatch, slack_dispatch, DISPATCH_MAP,
    )
    from src.simulator import WarehouseSimulator
    from src.features import FeatureExtractor, SCENARIO_FEATURE_NAMES

    sim_kw = {
        "base_arrival_rate":    config.get("base_arrival_rate", 2.5),
        "breakdown_prob":       config.get("breakdown_prob", 0.003),
        "batch_arrival_size":   config.get("batch_arrival_size", 30),
        "lunch_penalty_factor": config.get("lunch_penalty_factor", 1.3),
        "job_type_frequencies": config.get("job_type_frequencies", {}),
        "due_date_tightness":   config.get("due_date_tightness", 1.0),
        "processing_time_scale": config.get("processing_time_scale", 1.0),
    }

    seed = config["seed"]
    fe = FeatureExtractor()
    sim = WarehouseSimulator(seed=seed, heuristic_fn=fifo_dispatch, feature_extractor=fe, **sim_kw)
    sim.init()

    rows = []
    SIM_DURATION = 600.0

    for t in np.arange(SNAPSHOT_INTERVAL, SIM_DURATION, SNAPSHOT_INTERVAL):
        t = float(t)
        sim.step_to(t)
        state_snap = sim.get_state_snapshot()

        # Extract 22 scenario features from current state
        features = fe.extract_scenario_features(state_snap)
        if np.any(~np.isfinite(features)):
            continue  # skip bad windows

        # Save state for forking
        saved_state = sim.save_state()

        # Fork 6 heuristics for FORK_WINDOW min each, collect raw metrics
        fork_end = t + FORK_WINDOW
        raw_metrics: List[Tuple[float, float, float]] = []
        for heur_name in HEURISTIC_NAMES:
            try:
                heur_fn = DISPATCH_MAP[heur_name]
                fork = WarehouseSimulator.from_state(saved_state, heur_fn)
                fork.step_to(fork_end)
                metrics = fork.get_partial_metrics(since_time=t)
                tard = metrics.total_tardiness if np.isfinite(metrics.total_tardiness) else 1e9
                sla  = metrics.sla_breach_rate if np.isfinite(metrics.sla_breach_rate) else 1.0
                cyc  = metrics.avg_cycle_time if np.isfinite(metrics.avg_cycle_time) else 1e6
            except Exception:
                tard, sla, cyc = 1e9, 1.0, 1e6
            raw_metrics.append((tard, sla, cyc))

        # Normalize each metric across the 6 heuristics so units are comparable.
        # Without this, raw tardiness (hundreds-thousands) dominates SLA (0-1) and
        # cycle time (tens), so WSPT gets labeled at almost every snapshot.
        arr = np.array(raw_metrics, dtype=float)
        def _norm(col: np.ndarray) -> np.ndarray:
            lo, hi = float(col.min()), float(col.max())
            if hi - lo < 1e-10:
                return np.zeros_like(col)
            return (col - lo) / (hi - lo)
        n_tard = _norm(arr[:, 0])
        n_sla  = _norm(arr[:, 1])
        n_cyc  = _norm(arr[:, 2])
        # Weights match the benchmark objective (tardiness-dominant) to avoid
        # cycle-time over-weighting which biased labels toward WSPT.
        scores_arr = 0.55 * n_tard + 0.35 * n_sla + 0.10 * n_cyc

        # Label: best heuristic for THIS situation (lowest normalized composite).
        # Tie-break: when the top two are within TIE_EPS, break ties by the
        # heuristic that currently has the lower global label frequency.
        # This prevents any rule collapsing the dataset (WSPT dominance).
        TIE_EPS = 0.02
        order = np.argsort(scores_arr)
        best = int(order[0])
        runner = int(order[1]) if len(order) > 1 else best
        if abs(scores_arr[best] - scores_arr[runner]) < TIE_EPS:
            # Use rarity-of-label heuristic: among tied candidates, prefer the one
            # with lower ordinal frequency (approximated by reverse index order —
            # FIFO=0, EDD=1, CR=2, ATC=3, WSPT=4, Slack=5; non-WSPT preferred
            # when roughly equal).
            tied = [int(i) for i in order if scores_arr[i] - scores_arr[best] < TIE_EPS]
            # Prefer the tied heuristic furthest from WSPT (index 4) to diversify
            tied.sort(key=lambda h: abs(h - 4), reverse=True)
            best = tied[0]
        label = best
        scores = scores_arr.tolist()

        row = {name: float(val) for name, val in zip(SCENARIO_FEATURE_NAMES, features)}
        row["label"] = label
        rows.append(row)

    return rows


def _composite_score(metrics) -> float:
    """Scoring formula: 0.40*tardiness + 0.35*sla + 0.25*cycle_time (normalized)."""
    # Raw (unnormalized) — normalization happens across heuristics in the caller
    tard = metrics.total_tardiness if metrics.total_tardiness != float("inf") else 1e9
    sla = metrics.sla_breach_rate if metrics.sla_breach_rate != float("inf") else 1.0
    cyc = metrics.avg_cycle_time if metrics.avg_cycle_time != float("inf") else 1e6
    return 0.40 * tard + 0.35 * sla * 1000 + 0.25 * cyc


# ---------------------------------------------------------------------------
# Priority dataset worker (ported from DAHS_1)
# ---------------------------------------------------------------------------

def _run_priority_scenario(args: Tuple[int, int]) -> List[Dict[str, Any]]:
    """Worker: run one seed with ATC baseline, collect job-level feature rows."""
    seed, n_points = args
    from src.heuristics import atc_dispatch
    from src.simulator import WarehouseSimulator
    from src.features import FeatureExtractor

    _PRIO_W = {"A": 2.0, "B": 1.5, "C": 1.0, "D": 0.8, "E": 3.0}
    _DD_OFFSET = {"A": 120, "B": 160, "C": 240, "D": 320, "E": 60}

    fe = FeatureExtractor()
    sim = WarehouseSimulator(seed=seed, heuristic_fn=atc_dispatch, feature_extractor=fe)
    sim.run(duration=600.0)

    rows: List[Dict[str, Any]] = []
    state = sim.get_state_snapshot()
    completed = sim.completed_jobs

    if not completed:
        return rows

    rng = np.random.default_rng(seed)
    sampled = rng.choice(len(completed), size=min(n_points, len(completed)), replace=False)

    for idx in sampled:
        job = completed[int(idx)]
        scenario_feats = fe.extract_scenario_features(state)
        job_feats = fe.extract_job_features(job, state)

        w = _PRIO_W.get(job.job_type, 1.0)
        dd_off = _DD_OFFSET.get(job.job_type, 120)
        cycle_time = job.completion_time - job.arrival_time
        tardiness = max(0.0, job.completion_time - job.due_date)

        remaining = job.remaining_proc_time()
        time_to_due = job.due_date - state["current_time"]
        urgency = 1.0 - min(1.0, max(0.0, time_to_due / max(dd_off, 1.0)))
        importance = w / 3.0
        efficiency = 1.0 / (1.0 + remaining / 30.0)
        delivery_perf = max(0.0, 1.0 - tardiness / max(dd_off, 1.0))

        priority_score = float(
            0.30 * urgency
            + 0.25 * importance
            + 0.20 * efficiency
            + 0.25 * delivery_perf
        )

        if not np.isfinite(priority_score):
            continue

        row = {
            **{f"sf_{i}": float(v) for i, v in enumerate(scenario_feats)},
            **{f"jf_{i}": float(v) for i, v in enumerate(job_feats)},
            "priority_score": priority_score,
        }
        rows.append(row)

    return rows


# ---------------------------------------------------------------------------
# Dataset generators
# ---------------------------------------------------------------------------

def generate_selector_dataset(
    n_scenarios: int = 1000,
    n_workers: int = 4,
    save: bool = True,
) -> pd.DataFrame:
    """Generate the heuristic selector training dataset using snapshot-fork algorithm.

    Parameters
    ----------
    n_scenarios : int
        Number of scenario seeds to simulate.
    n_workers : int
        Number of parallel worker processes.
    save : bool
        Whether to save the CSV to data/raw/.

    Returns
    -------
    pd.DataFrame
        22 scenario feature columns + "label" (0-5, one per heuristic).
        ~60 rows per scenario (one per 10-min snapshot).
    """
    from src.features import SCENARIO_FEATURE_NAMES

    master_rng = np.random.default_rng(777)
    configs = _make_diverse_scenario_configs(n_scenarios, master_rng)

    logger.info(
        "Generating selector dataset (snapshot-fork): %d scenarios × ~60 snapshots each",
        n_scenarios
    )

    all_rows: List[Dict[str, Any]] = []
    ctx = mp.get_context("spawn")
    with ctx.Pool(processes=n_workers) as pool:
        for rows in tqdm(
            pool.imap_unordered(_run_snapshot_scenario, configs),
            total=len(configs),
            desc="Snapshot-fork data gen",
        ):
            all_rows.extend(rows)

    df = pd.DataFrame(all_rows)

    # Sanitize
    df = df.replace([np.inf, -np.inf], np.nan).fillna(0.0)

    logger.info("Selector dataset shape: %s", df.shape)
    if "label" in df.columns:
        label_counts = df["label"].value_counts().to_dict()
        logger.info("Label distribution: %s", label_counts)

    if save:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        path = DATA_DIR / "selector_dataset.csv"
        df.to_csv(path, index=False)
        logger.info("Saved selector dataset -> %s", path)

    return df


def generate_priority_dataset(
    n_scenarios: int = 5_000,
    n_points_per: int = 10,
    n_workers: int = 4,
    save: bool = True,
) -> pd.DataFrame:
    """Generate the priority predictor training dataset (ported from DAHS_1)."""
    from src.features import SCENARIO_FEATURE_NAMES, JOB_FEATURE_NAMES

    seeds = list(range(20_000, 20_000 + n_scenarios))
    all_args = [(seed, n_points_per) for seed in seeds]

    logger.info("Generating priority dataset: %d scenarios × %d points", n_scenarios, n_points_per)

    all_rows: List[Dict] = []
    ctx = mp.get_context("spawn")
    with ctx.Pool(processes=n_workers) as pool:
        for batch in tqdm(
            pool.imap_unordered(_run_priority_scenario, all_args),
            total=len(all_args),
            desc="Priority data gen",
        ):
            all_rows.extend(batch)

    df = pd.DataFrame(all_rows)
    df = df.replace([np.inf, -np.inf], np.nan).dropna()

    sf_names = {f"sf_{i}": name for i, name in enumerate(SCENARIO_FEATURE_NAMES)}
    jf_names = {f"jf_{i}": name for i, name in enumerate(JOB_FEATURE_NAMES)}
    df.rename(columns={**sf_names, **jf_names}, inplace=True)

    logger.info("Priority dataset shape: %s", df.shape)

    if save:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        path = DATA_DIR / "priority_dataset.csv"
        df.to_csv(path, index=False)
        logger.info("Saved priority dataset -> %s", path)

    return df


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    generate_selector_dataset(n_scenarios=50, n_workers=2)