# Prepared by Molla Samser (RSK World) | hello@rskworld.in | support@rskworld.in | +91 93305 39277 | Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India 713147 | https://rskworld.in
"""
Lightweight preprocessing pipeline for the Sales Forecasting dataset.

Steps:
- Load raw CSV
- Parse dates and sort
- Build calendar, lag, and rolling features
- Save processed CSV for modeling
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd


def add_time_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["date"] = pd.to_datetime(df["date"])
    df = df.sort_values(["product_id", "date"])

    df["year"] = df["date"].dt.year
    df["month"] = df["date"].dt.month
    df["week"] = df["date"].dt.isocalendar().week.astype(int)
    df["quarter"] = df["date"].dt.quarter
    df["day"] = df["date"].dt.day
    df["dayofweek"] = df["date"].dt.dayofweek
    df["is_weekend"] = df["dayofweek"].isin([5, 6]).astype(int)
    df["dayofyear"] = df["date"].dt.dayofyear
    df["month_start"] = df["date"].dt.is_month_start.astype(int)
    df["month_end"] = df["date"].dt.is_month_end.astype(int)
    df["trend_index"] = df.groupby("product_id").cumcount()

    # Seasonal Fourier terms (capture yearly cycles)
    for k in (1, 2):
        df[f"sin_dayofyear_{k}"] = np.sin(2 * np.pi * k * df["dayofyear"] / 365.25)
        df[f"cos_dayofyear_{k}"] = np.cos(2 * np.pi * k * df["dayofyear"] / 365.25)

    # Ensure revenue column exists even if upstream omitted it.
    if "revenue" not in df.columns:
        df["revenue"] = df["price"] * df["units_sold"]

    # Group-wise lag and rolling stats per product_id
    group_cols = ["product_id"]
    for lag in (7, 14, 28):
        df[f"lag_{lag}_units"] = df.groupby(group_cols)["units_sold"].shift(lag // 7)
        df[f"lag_{lag}_revenue"] = df.groupby(group_cols)["revenue"].shift(lag // 7)

    # Rolling windows expressed in weeks (4 ~ monthly, 12 ~ quarterly)
    df["rolling_4w_units_mean"] = (
        df.groupby(group_cols)["units_sold"].transform(lambda s: s.rolling(window=4, min_periods=1).mean())
    )
    df["rolling_4w_revenue_sum"] = (
        df.groupby(group_cols)["revenue"].transform(lambda s: s.rolling(window=4, min_periods=1).sum())
    )
    df["rolling_12w_units_mean"] = (
        df.groupby(group_cols)["units_sold"].transform(lambda s: s.rolling(window=12, min_periods=1).mean())
    )
    df["rolling_12w_revenue_sum"] = (
        df.groupby(group_cols)["revenue"].transform(lambda s: s.rolling(window=12, min_periods=1).sum())
    )

    # Promotion indicator
    if "promotion" in df.columns:
        df["has_promo"] = df["promotion"].fillna("").str.strip().ne("").astype(int)

    # Price dynamics
    if "price" in df.columns:
        df["price_change_pct"] = (
            df.groupby(group_cols)["price"].pct_change().replace([np.inf, -np.inf], np.nan).fillna(0.0)
        )

    # Revenue per unit (guard against divide-by-zero)
    df["revenue_per_unit"] = df["revenue"] / df["units_sold"].replace({0: np.nan})
    df["revenue_per_unit"] = df["revenue_per_unit"].fillna(0.0)

    # Inventory coverage in days: stock divided by average daily sales (from 4w mean)
    if "stock_level" in df.columns:
        avg_daily = df["rolling_4w_units_mean"] / 7
        df["inventory_coverage_days"] = (df["stock_level"] / avg_daily.replace(0, np.nan)).fillna(0.0)

    return df


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Preprocess sales dataset for forecasting")
    parser.add_argument("--input", type=Path, default=Path("data") / "sample_sales.csv", help="Input CSV path")
    parser.add_argument("--output", type=Path, default=Path("data") / "processed_sales.csv", help="Output CSV path")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    df = pd.read_csv(args.input)
    processed = add_time_features(df)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    processed.to_csv(args.output, index=False)
    print(f"Saved processed dataset to {args.output}")


if __name__ == "__main__":
    main()