Grid_DataLoader.py

from pathlib import Path

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset


# ---------------------------------------------------------------------
# 1. Sparse Dataset 정의 (panel + static → ST-ResNet 입력 생성)
# ---------------------------------------------------------------------
class GridSTSparseDataset(Dataset):
    """
    ST-ResNet 계열 모델용 시공간 Grid 데이터셋.

    동적(dynamic) 피처 (F_dyn = 3):
        1) cnt(g,t)    : 사고 건수
        2) dist(g,t)   : 평균 현장거리 (mean_dist, meter)
        3) traffic(g,t): GRID_SPEED (평균 교통량/속도)

      → 한 시점 동적 텐서 shape: (H, W, 3)
        예) closeness 입력: (L_c, H, W, 3)

    정적(static) 피처 (F_stat = 3):
        1) road_len_in_grid(g)
        2) n_links(g)
        3) n_segments(g)

      → 정적 텐서 shape: (H, W, 3)

    반환 dict:
        {
            "X_closeness": (L_c, H, W, 3),
            "X_period":    (L_p, H, W, 3) or None,
            "X_trend":     (L_t, H, W, 3) or None,
            "X_static":    (H, W, 3),
            "y":           (H, W),         # 평균 소요시간
            "mask":        (H, W),         # 관측 마스크
            "t":           int,            # t_idx
            "ext":         (3,) or None,   # [hour/23, dow/6, is_weekend]
        }
    """

    def __init__(
        self,
        panel: pd.DataFrame,
        static: np.ndarray,
        T: int,
        R: int,
        C: int,
        seq_len: int = 6,
        period: int = 24,
        trend: int = 24 * 7,
        num_periods: int = 3,
        num_trends: int = 2,
        hour_arr: np.ndarray | None = None,
        dow_arr: np.ndarray | None = None,
        is_weekend_arr: np.ndarray | None = None,
    ):
        """
        panel 컬럼 요구사항:
            - "t_idx"
            - "r_idx", "c_idx"
            - "cnt"          : 사고 건수
            - "mean_time"    : 평균 소요시간 (초)
            - "mean_dist"    : 평균 현장거리
            - "traffic"      : GRID_SPEED (평균 속도/교통량)
        """
        super().__init__()
        self.panel = panel
        self.static = torch.from_numpy(static).float()  # (R, C, F_stat=3)
        self.T = T
        self.R = R
        self.C = C

        self.seq_len = seq_len
        self.period = period
        self.trend = trend
        self.num_periods = num_periods
        self.num_trends = num_trends

        self.hour_arr = hour_arr
        self.dow_arr = dow_arr
        self.is_weekend_arr = is_weekend_arr

        # (t_idx → sparse grid 정보) 캐시
        # time_dict[t] = (r_idx, c_idx, cnt, mean_time, mean_dist, traffic)
        self._build_time_index()

        # temporal dependency 만족하는 t만 유효
        min_t = seq_len
        if period > 0 and num_periods > 0:
            min_t = max(min_t, num_periods * period)
        if trend > 0 and num_trends > 0:
            min_t = max(min_t, num_trends * trend)

        self.min_t = min_t
        self.valid_t = list(range(self.min_t, self.T))

        print(
            f"[GridSTSparseDataset] T={self.T}, R={self.R}, C={self.C}, "
            f"seq_len={self.seq_len}, period={self.period}, trend={self.trend}, "
            f"|valid_t|={len(self.valid_t)}"
        )

    def _build_time_index(self):
        """
        panel를 t_idx별로 그룹핑해서 sparse grid 저장.

        time_dict[t] = (
            r_idx: int[],
            c_idx: int[],
            cnt: float32[],
            mean_time: float32[],
            mean_dist: float32[],
            traffic: float32[],
        )
        """
        self.time_dict = {}
        g = self.panel.groupby("t_idx")
        for t, df_t in g:
            r = df_t["r_idx"].to_numpy(dtype=int)
            c = df_t["c_idx"].to_numpy(dtype=int)

            cnt = df_t["cnt"].to_numpy(dtype=np.float32)
            mt = df_t["mean_time"].to_numpy(dtype=np.float32)
            dist = df_t["mean_dist"].to_numpy(dtype=np.float32)
            traf = df_t["traffic"].to_numpy(dtype=np.float32)

            self.time_dict[int(t)] = (r, c, cnt, mt, dist, traf)

    def __len__(self):
        return len(self.valid_t)

    def _build_map_at_t(self, t: int):
        """
        시점 t에 대해 full grid 맵 생성.

        Returns
        -------
        cnt_map   : (R, C), 사고 건수
        dist_map  : (R, C), 평균 현장거리
        traf_map  : (R, C), 평균 교통량/속도
        y_map     : (R, C), 평균 소요시간 (초)
        mask_map  : (R, C), 관측 마스크 (0/1)
        """
        R, C = self.R, self.C

        cnt_map = np.zeros((R, C), dtype=np.float32)
        dist_map = np.zeros((R, C), dtype=np.float32)
        traf_map = np.zeros((R, C), dtype=np.float32)
        y_map = np.zeros((R, C), dtype=np.float32)

        if t in self.time_dict:
            r_idx, c_idx, cnt, mt, dist, traf = self.time_dict[t]

            cnt_map[r_idx, c_idx] = cnt
            dist_map[r_idx, c_idx] = dist
            traf_map[r_idx, c_idx] = traf
            y_map[r_idx, c_idx] = mt

        # 0 < y < 3600 인 지점만 유효 관측
        mask_map = ((y_map > 0) & (y_map < 3600)).astype(np.float32)

        return cnt_map, dist_map, traf_map, y_map, mask_map

    def __getitem__(self, idx):
        t = self.valid_t[idx]

        # ---------- closeness (L_c 시점, F_dyn=3) ----------
        Xc_cnt = []
        Xc_dist = []
        Xc_traf = []

        for k in range(self.seq_len, 0, -1):
            t_k = t - k
            cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k)
            Xc_cnt.append(cnt_map)
            Xc_dist.append(dist_map)
            Xc_traf.append(traf_map)

        # (L_c, R, C, 3) = [cnt, dist, traffic]
        X_c = np.stack(
            [
                np.stack(Xc_cnt, axis=0),
                np.stack(Xc_dist, axis=0),
                np.stack(Xc_traf, axis=0),
            ],
            axis=-1,
        )

        # ---------- period ----------
        if self.period > 0 and self.num_periods > 0:
            Xp_cnt = []
            Xp_dist = []
            Xp_traf = []

            for k in range(1, self.num_periods + 1):
                t_k = t - k * self.period
                cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k)
                Xp_cnt.append(cnt_map)
                Xp_dist.append(dist_map)
                Xp_traf.append(traf_map)

            # (L_p, R, C, 3)
            X_p = np.stack(
                [
                    np.stack(Xp_cnt, axis=0),
                    np.stack(Xp_dist, axis=0),
                    np.stack(Xp_traf, axis=0),
                ],
                axis=-1,
            )
        else:
            X_p = None

        # ---------- trend ----------
        if self.trend > 0 and self.num_trends > 0:
            Xt_cnt = []
            Xt_dist = []
            Xt_traf = []

            for k in range(1, self.num_trends + 1):
                t_k = t - k * self.trend
                cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k)
                Xt_cnt.append(cnt_map)
                Xt_dist.append(dist_map)
                Xt_traf.append(traf_map)

            # (L_t, R, C, 3)
            X_t = np.stack(
                [
                    np.stack(Xt_cnt, axis=0),
                    np.stack(Xt_dist, axis=0),
                    np.stack(Xt_traf, axis=0),
                ],
                axis=-1,
            )
        else:
            X_t = None

        # ---------- target & mask ----------
        _, _, _, y_map, mask_map = self._build_map_at_t(t)

        # ---------- static ----------
        X_s = self.static  # (R, C, 3): [road_len_in_grid, n_links, n_segments]

        # ---------- 외생 변수 ----------
        if (
            self.hour_arr is not None
            and self.dow_arr is not None
            and self.is_weekend_arr is not None
        ):
            h = float(self.hour_arr[t]) / 23.0
            d = float(self.dow_arr[t]) / 6.0
            w = float(self.is_weekend_arr[t])
            ext = torch.tensor([h, d, w], dtype=torch.float32)
        else:
            ext = None

        # Torch 변환
        X_c = torch.from_numpy(X_c).float()       # (L_c, R, C, 3)
        y_t = torch.from_numpy(y_map).float()     # (R, C)
        m_t = torch.from_numpy(mask_map).float()  # (R, C)
        X_s = X_s.float()                         # (R, C, 3)

        if X_p is not None:
            X_p = torch.from_numpy(X_p).float()   # (L_p, R, C, 3)
        if X_t is not None:
            X_t = torch.from_numpy(X_t).float()   # (L_t, R, C, 3)

        return {
            "X_closeness": X_c,
            "X_period": X_p,
            "X_trend": X_t,
            "X_static": X_s,
            "y": y_t,
            "mask": m_t,
            "t": t,
            "ext": ext,
        }

# ---------------------------------------------------------------------
# 2. 해상도별 parquet → panel / static / 시간메타 생성
# ---------------------------------------------------------------------
def build_sparse_panel_and_static(GRID_RES_M: int):
    """
    준비된 parquet에서 모델 입력용 panel / static / 시간 메타데이터 생성.

    동적(panel) (t_idx, r_idx, c_idx):
        - cnt        : cnt_events (사고 건수)
        - mean_time  : 소요시각 평균
        - mean_dist  : 현장거리 평균
        - traffic    : GRID_SPEED (평균 교통량/속도)

    정적(static) (R, C, 3):
        - road_len_in_grid
        - n_links
        - n_segments
    """
    BASE_DIR = Path(__file__).resolve().parent

    DB_DIR = BASE_DIR / "DB"
    PREP_DIR = DB_DIR / "prepared"

    PATH_FACT_ACC_GRID = PREP_DIR / f"fact_accident_dispatch_grid_{GRID_RES_M}m.parquet"
    PATH_DYN_FEAT      = PREP_DIR / f"dyn_features_grid_{GRID_RES_M}m.parquet"
    PATH_GRID_STATIC   = PREP_DIR / f"grid_static_road_{GRID_RES_M}m.parquet"
    PATH_TIME_BIAS     = PREP_DIR / "time_bias_features.parquet"

    print(f"[{GRID_RES_M}m] Loading prepared parquet files...")

    fact = pd.read_parquet(PATH_FACT_ACC_GRID)
    dyn = pd.read_parquet(PATH_DYN_FEAT)
    grid_static = pd.read_parquet(PATH_GRID_STATIC)
    time_bias = pd.read_parquet(PATH_TIME_BIAS)

    print(f"[{GRID_RES_M}m] fact shape:", fact.shape)
    print(f"[{GRID_RES_M}m] dyn shape:", dyn.shape)
    print(f"[{GRID_RES_M}m] grid_static shape:", grid_static.shape)
    print(f"[{GRID_RES_M}m] time_bias shape:", time_bias.shape)

    # r_idx / c_idx 준비
    if "r_idx" in fact.columns and "c_idx" in fact.columns:
        print(f"[{GRID_RES_M}m] Using existing r_idx / c_idx from fact table.")
    else:
        print(f"[{GRID_RES_M}m] r_idx / c_idx 없음. grid_row / grid_col 기준으로 생성.")
        row_vals = np.sort(fact["grid_row"].unique())
        col_vals = np.sort(fact["grid_col"].unique())
        row2idx = {r: i for i, r in enumerate(row_vals)}
        col2idx = {c: i for i, c in enumerate(col_vals)}
        fact["r_idx"] = fact["grid_row"].map(row2idx)
        fact["c_idx"] = fact["grid_col"].map(col2idx)

    T = int(fact["t_idx"].max()) + 1
    R = int(fact["r_idx"].max()) + 1
    C = int(fact["c_idx"].max()) + 1
    print(f"[{GRID_RES_M}m] T (hours) = {T}, R = {R}, C = {C}")

    # ---------------- 시간 메타데이터 (time_bias) ----------------
    time_bias = time_bias.sort_values("t_idx").drop_duplicates(subset=["t_idx"])
    if time_bias["t_idx"].max() + 1 < T:
        T = int(time_bias["t_idx"].max()) + 1
        print(f"[{GRID_RES_M}m] Adjusted T to {T} based on time_bias.")

    hour_arr = np.zeros(T, dtype=np.int16)
    dow_arr = np.zeros(T, dtype=np.int16)
    is_weekend_arr = np.zeros(T, dtype=np.int8)

    t_idx_meta = time_bias["t_idx"].to_numpy(dtype=int)
    hours = time_bias["hour"].to_numpy(dtype=int)
    dows = time_bias["weekday"].to_numpy(dtype=int)
    weekends = time_bias["is_weekend"].to_numpy(dtype=np.int8)

    hour_arr[t_idx_meta] = hours
    dow_arr[t_idx_meta] = dows
    is_weekend_arr[t_idx_meta] = weekends

    # ---------------- panel 생성 ----------------
    print(f"[{GRID_RES_M}m] Aggregating mean response time per (t_idx, r_idx, c_idx)...")

    if "소요시각" not in fact.columns:
        raise ValueError(f"[{GRID_RES_M}m] fact 테이블에 '소요시각' 컬럼이 없습니다.")

    # (t_idx, r_idx, c_idx)별 평균 소요시각
    agg_y = (
        fact
        .groupby(["t_idx", "r_idx", "c_idx"])
        .agg(
            mean_time=("소요시각", "mean"),
        )
        .reset_index()
    )

    # dyn: (grid_id, 단위_m, t_idx)별 cnt_events, mean_dist, GRID_SPEED, link_cnt
    grid_map = (
        fact[["grid_id", "단위_m", "r_idx", "c_idx"]]
        .drop_duplicates(subset=["grid_id", "단위_m"])
    )

    dyn = dyn.merge(
        grid_map,
        on=["grid_id", "단위_m"],
        how="left",
        validate="many_to_one",
    )

    dyn = dyn.dropna(subset=["r_idx", "c_idx"]).copy()
    dyn["r_idx"] = dyn["r_idx"].astype(int)
    dyn["c_idx"] = dyn["c_idx"].astype(int)

    # dyn + mean_time merge
    panel = dyn.merge(
        agg_y,
        on=["t_idx", "r_idx", "c_idx"],
        how="left",
        validate="one_to_one",
    )

    # 이름 정리: 모델에서 사용할 이름
    panel = panel.rename(
        columns={
            "cnt_events": "cnt",
            "GRID_SPEED": "traffic",   # 동적 traffic 피처
            # link_cnt는 지금은 사용 안 함 (원하면 static으로 따로 쓸 수 있음)
        }
    )

    for col in ["cnt", "mean_dist", "traffic", "mean_time"]:
        if col not in panel.columns:
            raise ValueError(f"[{GRID_RES_M}m] panel에 '{col}' 컬럼이 없습니다.")
        panel[col] = panel[col].fillna(0.0).astype(np.float32)

    panel = panel[["t_idx", "r_idx", "c_idx",
                   "cnt", "mean_time", "mean_dist", "traffic"]].copy()

    print(f"[{GRID_RES_M}m] panel rows:", len(panel))

    # ---------------- static: (R, C, 3) ----------------
    print(f"[{GRID_RES_M}m] Building static grid features...")

    row_map = fact[["grid_row", "r_idx"]].drop_duplicates()
    col_map = fact[["grid_col", "c_idx"]].drop_duplicates()
    row2idx = dict(zip(row_map["grid_row"], row_map["r_idx"]))
    col2idx = dict(zip(col_map["grid_col"], col_map["c_idx"]))

    grid_static["r_idx"] = grid_static["grid_row"].map(row2idx)
    grid_static["c_idx"] = grid_static["grid_col"].map(col2idx)

    grid_static = grid_static.dropna(subset=["r_idx", "c_idx"]).copy()
    grid_static["r_idx"] = grid_static["r_idx"].astype(int)
    grid_static["c_idx"] = grid_static["c_idx"].astype(int)

    for col in ["road_len_in_grid", "n_links", "n_segments"]:
        if col not in grid_static.columns:
            grid_static[col] = 0.0
        else:
            grid_static[col] = grid_static[col].fillna(0.0)

    features = ["road_len_in_grid", "n_links", "n_segments"]
    F_stat = len(features)
    static = np.zeros((R, C, F_stat), dtype=np.float32)

    r_idx_arr = grid_static["r_idx"].to_numpy(dtype=int)
    c_idx_arr = grid_static["c_idx"].to_numpy(dtype=int)

    for i, feat in enumerate(features):
        vals = grid_static[feat].to_numpy(dtype=np.float32)
        static[r_idx_arr, c_idx_arr, i] = vals

    print(f"[{GRID_RES_M}m] static shape:", static.shape)

    return panel, static, T, R, C, hour_arr, dow_arr, is_weekend_arr