"""Build the RestockIQ feature table from the three manually-downloaded M5 CSVs.

Reads (fixed paths, no Kaggle API):
    data/raw/calendar.csv
    data/raw/sell_prices.csv
    data/raw/sales_train_evaluation.csv

Writes:
    data/features.parquet     long feature table for LightGBM training
    data/sales_wide.npy       (n_series, 1941) float32 sales matrix for recursive forecasting
    data/series_meta.parquet  one row per series: id, item_id, dept_id, cat_id, store_id, state_id
    data/day_map.parquet      one row per d: date, wm_yr_wk, dow, month, snap flags, event codes

The melted long table is ~59M rows; dtypes are downcast aggressively (int8/int16/float32,
categoricals) immediately after melting so the whole thing fits in RAM.
"""

import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from app import config as cfg

MISSING_DATA_MSG = """
ERROR: missing M5 input file(s):
{missing}

These must be downloaded manually (no API token needed) — see README section
"Manual setup" (spec Section 10c):
  1. Go to https://www.kaggle.com/competitions/m5-forecasting-accuracy/data
  2. Log in, accept the competition rules, click "Download All".
  3. Unzip and place these three files at:
       data/raw/calendar.csv
       data/raw/sell_prices.csv
       data/raw/sales_train_evaluation.csv
"""


def check_raw_files() -> None:
    missing = [str(p) for p in (cfg.RAW_CALENDAR, cfg.RAW_PRICES, cfg.RAW_SALES) if not p.exists()]
    if missing:
        print(MISSING_DATA_MSG.format(missing="\n".join("  - " + m for m in missing)))
        sys.exit(1)


def rolling_from_cumsum(vals: np.ndarray, window: int) -> np.ndarray:
    """Mean of the `window` days strictly before each day t (t-window .. t-1), per row.

    vals: (n_series, n_days). Returns float32 with NaN where the window is incomplete.
    """
    n, t = vals.shape
    csum = np.zeros((n, t + 1), dtype=np.float64)
    np.cumsum(vals, axis=1, dtype=np.float64, out=csum[:, 1:])
    out = np.full((n, t), np.nan, dtype=np.float32)
    out[:, window:] = ((csum[:, window:-1] - csum[:, :-window - 1]) / window).astype(np.float32)
    return out


def rolling_std_from_cumsum(vals: np.ndarray, window: int) -> np.ndarray:
    """Std (population) of the `window` days strictly before each day t, per row."""
    n, t = vals.shape
    csum = np.zeros((n, t + 1), dtype=np.float64)
    np.cumsum(vals, axis=1, dtype=np.float64, out=csum[:, 1:])
    csum2 = np.zeros((n, t + 1), dtype=np.float64)
    np.cumsum(vals.astype(np.float64) ** 2, axis=1, dtype=np.float64, out=csum2[:, 1:])
    out = np.full((n, t), np.nan, dtype=np.float32)
    s = csum[:, window:-1] - csum[:, :-window - 1]
    s2 = csum2[:, window:-1] - csum2[:, :-window - 1]
    var = s2 / window - (s / window) ** 2
    out[:, window:] = np.sqrt(np.clip(var, 0.0, None)).astype(np.float32)
    return out


def lag(vals: np.ndarray, k: int) -> np.ndarray:
    """Value k days earlier, per row. NaN where unavailable."""
    n, t = vals.shape
    out = np.full((n, t), np.nan, dtype=np.float32)
    out[:, k:] = vals[:, :-k]
    return out


def build_day_map(calendar: pd.DataFrame) -> pd.DataFrame:
    """One row per d_1..d_1941 with calendar-derived per-day features."""
    cal = calendar.copy()
    cal["d_idx"] = cal["d"].str.slice(2).astype(np.int16)
    cal = cal[cal["d_idx"] <= cfg.N_DAYS].sort_values("d_idx").reset_index(drop=True)
    day_map = pd.DataFrame(
        {
            "d": cal["d_idx"],
            "date": pd.to_datetime(cal["date"]),
            "wm_yr_wk": cal["wm_yr_wk"].astype(np.int32),
            # wday: 1=Saturday..7=Friday (M5 convention); keep as-is, it's just a category
            "day_of_week": cal["wday"].astype(np.int8),
            "month": cal["month"].astype(np.int8),
            "snap_CA": cal["snap_CA"].astype(np.int8),
            "snap_TX": cal["snap_TX"].astype(np.int8),
            "snap_WI": cal["snap_WI"].astype(np.int8),
            "event_name": cal["event_name_1"].fillna("none").astype("category"),
            "event_type": cal["event_type_1"].fillna("none").astype("category"),
        }
    )
    assert len(day_map) == cfg.N_DAYS, f"expected {cfg.N_DAYS} days, got {len(day_map)}"
    return day_map


def build_price_matrix(
    prices: pd.DataFrame, series_meta: pd.DataFrame, week_ids: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """(n_series, n_weeks) price matrix + int8 price-change flag matrix.

    NaN price = product not offered that week (M5 convention: before first sale).
    """
    n_series = len(series_meta)
    week_pos = {int(w): i for i, w in enumerate(week_ids)}
    n_weeks = len(week_ids)

    # calendar extends past d_1941, so sell_prices has weeks beyond our horizon: drop them
    prices = prices[prices["wm_yr_wk"].isin(week_pos.keys())]

    key_to_row = pd.Series(
        np.arange(n_series, dtype=np.int32),
        index=series_meta["store_id"].astype(str) + "|" + series_meta["item_id"].astype(str),
    )
    rows = key_to_row.reindex(prices["store_id"] + "|" + prices["item_id"]).to_numpy()
    cols = prices["wm_yr_wk"].map(week_pos).to_numpy()

    price_mat = np.full((n_series, n_weeks), np.nan, dtype=np.float32)
    valid = ~pd.isna(rows) & ~pd.isna(cols)
    price_mat[rows[valid].astype(np.int64), cols[valid].astype(np.int64)] = (
        prices["sell_price"].to_numpy(dtype=np.float32)[valid]
    )

    # price changed vs previous offered week (first offered week counts as no change)
    prev = np.roll(price_mat, 1, axis=1)
    prev[:, 0] = np.nan
    change = ((price_mat != prev) & ~np.isnan(price_mat) & ~np.isnan(prev)).astype(np.int8)
    return price_mat, change


def main() -> None:
    check_raw_files()

    print("Reading calendar and prices ...")
    calendar = pd.read_csv(cfg.RAW_CALENDAR)
    prices = pd.read_csv(
        cfg.RAW_PRICES,
        dtype={"store_id": "string", "item_id": "string", "wm_yr_wk": np.int32,
               "sell_price": np.float32},
    )
    day_map = build_day_map(calendar)

    print("Reading wide sales table ...")
    d_cols = [f"d_{i}" for i in range(1, cfg.N_DAYS + 1)]
    sales = pd.read_csv(cfg.RAW_SALES, dtype={c: np.int16 for c in d_cols})
    n_series = len(sales)
    print(f"  {n_series} series x {cfg.N_DAYS} days")

    series_meta = sales[["id", "item_id", "dept_id", "cat_id", "store_id", "state_id"]].copy()
    for c in series_meta.columns:
        series_meta[c] = series_meta[c].astype("category")

    vals = sales[d_cols].to_numpy(dtype=np.float32)  # (n_series, n_days)
    del sales

    print("Computing lag / rolling features on the wide matrix ...")
    feats_wide = {f"lag_{k}": lag(vals, k) for k in cfg.LAGS}
    for w in cfg.ROLL_MEAN_WINDOWS:
        feats_wide[f"rolling_mean_{w}"] = rolling_from_cumsum(vals, w)
    feats_wide[f"rolling_std_{cfg.ROLL_STD_WINDOW}"] = rolling_std_from_cumsum(
        vals, cfg.ROLL_STD_WINDOW
    )

    print("Building price matrix ...")
    week_ids = day_map["wm_yr_wk"].unique()
    price_mat, price_change_mat = build_price_matrix(prices, series_meta, week_ids)
    week_idx_of_day = pd.Series(
        np.arange(len(week_ids)), index=week_ids
    )[day_map["wm_yr_wk"]].to_numpy()  # (n_days,)
    del prices

    print("Melting to long format (~59M rows) with downcast dtypes ...")
    n_days = cfg.N_DAYS
    series_rep = np.repeat(np.arange(n_series, dtype=np.int32), n_days)
    d_tile = np.tile(np.arange(1, n_days + 1, dtype=np.int16), n_series)

    long = pd.DataFrame({"d": d_tile, "sales": vals.ravel().astype(np.int16)})
    for name, mat in feats_wide.items():
        long[name] = mat.ravel()
        feats_wide[name] = None  # free as we go
    del feats_wide

    # id columns as categoricals built straight from codes (cheap, no string repeat)
    for c in ("id", "item_id", "dept_id", "cat_id", "store_id", "state_id"):
        col = series_meta[c]
        long[c] = pd.Categorical.from_codes(
            np.repeat(col.cat.codes.to_numpy(), n_days), categories=col.cat.categories
        )

    # per-day calendar features via day-index take
    day_idx = d_tile.astype(np.int32) - 1
    long["day_of_week"] = day_map["day_of_week"].to_numpy()[day_idx]
    long["month"] = day_map["month"].to_numpy()[day_idx]
    for c in ("event_name", "event_type"):
        long[c] = pd.Categorical.from_codes(
            day_map[c].cat.codes.to_numpy()[day_idx], categories=day_map[c].cat.categories
        )

    # snap flag for the series' own state
    snap_mat = day_map[["snap_CA", "snap_TX", "snap_WI"]].to_numpy()  # (n_days, 3)
    state_to_col = {s: i for i, s in enumerate(("CA", "TX", "WI"))}
    state_col_idx = series_meta["state_id"].map(state_to_col).to_numpy(dtype=np.int8)
    long["snap"] = snap_mat[day_idx, np.repeat(state_col_idx, n_days)].astype(np.int8)

    # price + price-change flag
    wk = week_idx_of_day[day_idx]
    long["sell_price"] = price_mat[series_rep, wk]
    long["price_change"] = price_change_mat[series_rep, wk]
    del price_mat, price_change_mat, series_rep, day_idx, wk

    before = len(long)
    # Drop rows where the product wasn't offered (no price -> no possible sale, M5 convention)
    long = long[long["sell_price"].notna()]
    print(f"  dropped {before - len(long):,} not-yet-offered rows -> {len(long):,} rows")

    print("Memory usage (MB):", int(long.memory_usage(deep=False).sum() / 1e6))

    cfg.DATA_DIR.mkdir(parents=True, exist_ok=True)
    print("Writing outputs ...")
    long.to_parquet(cfg.FEATURES_PARQUET, index=False)
    np.save(cfg.SALES_WIDE_NPY, vals)
    series_meta.to_parquet(cfg.SERIES_META_PARQUET, index=False)
    day_map.to_parquet(cfg.DAY_MAP_PARQUET, index=False)
    print(f"Done. features={cfg.FEATURES_PARQUET} rows={len(long):,}")


if __name__ == "__main__":
    main()