"""Feature engineering for the preparation-time model.

Highlights:
- ``kitchen_load`` is derived from order timestamps: number of other orders that
  fell into a 30-minute rolling window before each order, per city. The raw
  dataset has no explicit "kitchen load" feature; we approximate it through
  order density, which is the most direct proxy for "how busy was the kitchen
  when the order came in".
- Time-of-day buckets, weekday/weekend, distance.

The output dataframe is what ``train.py`` consumes.
"""

from __future__ import annotations

import sys
from math import asin, cos, radians, sin, sqrt
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.config import PROCESSED_DIR  # noqa: E402

CLEAN_CSV = PROCESSED_DIR / "food_delivery_clean.csv"
FEATURES_CSV = PROCESSED_DIR / "food_delivery_features.csv"

TARGET = "Time_taken(min)"

NUMERIC_FEATURES = [
    "Delivery_person_Age",
    "Delivery_person_Ratings",
    "Vehicle_condition",
    "multiple_deliveries",
    "distance_km",
    "distance_missing",
    "kitchen_load_30min",
    "hour_of_day",
    "day_of_week",
    "is_peak_hour",
    "is_weekend",
]

CATEGORICAL_FEATURES = [
    "Weatherconditions",
    "Road_traffic_density",
    "Type_of_order",
    "Type_of_vehicle",
    "Festival",
    "City",
]


def haversine_km(lat1: pd.Series, lon1: pd.Series, lat2: pd.Series, lon2: pd.Series) -> pd.Series:
    r = 6371.0
    lat1r = np.radians(lat1)
    lat2r = np.radians(lat2)
    dlat = np.radians(lat2 - lat1)
    dlon = np.radians(lon2 - lon1)
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlon / 2) ** 2
    return 2 * r * np.arcsin(np.sqrt(a))


def add_distance(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    if {
        "Restaurant_latitude",
        "Restaurant_longitude",
        "Delivery_location_latitude",
        "Delivery_location_longitude",
    }.issubset(df.columns):
        df["distance_km"] = haversine_km(
            df["Restaurant_latitude"],
            df["Restaurant_longitude"],
            df["Delivery_location_latitude"],
            df["Delivery_location_longitude"],
        )
        # cap absurd distances (data noise)
        df.loc[df["distance_km"] > 50, "distance_km"] = np.nan
    else:
        df["distance_km"] = np.nan
    return df


def add_time_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["order_datetime"] = pd.to_datetime(df["order_datetime"], errors="coerce")
    df["hour_of_day"] = df["order_datetime"].dt.hour
    df["day_of_week"] = df["order_datetime"].dt.dayofweek
    df["is_weekend"] = df["day_of_week"].isin([5, 6]).astype(int)
    df["is_peak_hour"] = df["hour_of_day"].isin([12, 13, 19, 20, 21]).astype(int)
    return df


def add_kitchen_load(df: pd.DataFrame, window_minutes: int = 30) -> pd.DataFrame:
    """Count of other orders within the previous ``window_minutes`` per city.

    The City column is a noisy categorical proxy for "restaurant location".
    We fall back to a single global grouping if the column is missing.
    """
    df = df.copy()
    df = df.sort_values("order_datetime")
    group_col = "City" if "City" in df.columns else None

    def _per_group(g: pd.DataFrame) -> pd.DataFrame:
        g = g.set_index("order_datetime").sort_index()
        rolling = g.assign(_one=1)["_one"].rolling(f"{window_minutes}min").sum() - 1
        g["kitchen_load_30min"] = rolling.values
        return g.reset_index()

    if group_col is not None and df[group_col].notna().any():
        df = df.groupby(group_col, group_keys=False).apply(_per_group)
    else:
        df = _per_group(df)

    df["kitchen_load_30min"] = df["kitchen_load_30min"].clip(lower=0).fillna(0)
    return df


def build_features(df: pd.DataFrame) -> pd.DataFrame:
    df = add_distance(df)
    # Explicit missing-distance flag so the model can learn the missing-data
    # pattern instead of relying on a silent median imputation. Captures both
    # missing coordinates and distances capped as noise (> 50 km).
    df["distance_missing"] = df["distance_km"].isna().astype(int)
    df = add_time_features(df)
    df = add_kitchen_load(df)

    keep = [TARGET] + NUMERIC_FEATURES + [c for c in CATEGORICAL_FEATURES if c in df.columns]
    keep = [c for c in keep if c in df.columns]
    out = df[keep].copy()

    # Drop rows with missing target
    out = out.dropna(subset=[TARGET])
    return out


def main() -> None:
    if not CLEAN_CSV.exists():
        raise FileNotFoundError(
            f"{CLEAN_CSV} missing. Run 'python -m src.ml.prepare_data' first."
        )
    df = pd.read_csv(CLEAN_CSV)
    feats = build_features(df)
    FEATURES_CSV.parent.mkdir(parents=True, exist_ok=True)
    feats.to_csv(FEATURES_CSV, index=False)
    print(f"[features] wrote {FEATURES_CSV}  ({len(feats):,} rows, {feats.shape[1]} cols)")


if __name__ == "__main__":
    main()