src/segmentation/clustering.py

"""
Product segmentation via demand-based feature engineering + HDBSCAN clustering.

Features per SKU:
  - Mean demand level
  - Coefficient of variation (volatility)
  - Event sensitivity (demand lift on event days)
  - Trend slope (linear regression over time)
  - Seasonality strength (std of monthly means / overall std)

Clusters are mapped to human-readable labels for stakeholder consumption.
"""

import numpy as np
import pandas as pd
import hdbscan
from sklearn.preprocessing import StandardScaler


# ── Segment label rules ───────────────────────────────────────────────────────
# Applied after clustering; labels are inferred from cluster centroid profiles.
SEGMENT_LABELS = {
    "high_volume_stable":    "Steady Staples",
    "high_volume_volatile":  "Event-Driven Heroes",
    "low_volume_volatile":   "Volatile Accessories",
    "declining":             "Declining SKUs",
    "growing":               "Rising Stars",
    "noise":                 "Unique / Unclassified",
}


def build_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Compute one row per unique_id with demand-based features.
    Returns a DataFrame indexed by unique_id.
    """
    df = df.copy()
    df["ds"] = pd.to_datetime(df["ds"])
    df["t"] = (df["ds"] - df["ds"].min()).dt.days

    features = {}

    for uid, group in df.groupby("unique_id"):
        group = group.sort_values("ds")
        y = group["y"].values
        t = group["t"].values

        # Mean demand
        mean_demand = y.mean()

        # Coefficient of variation
        cv = y.std() / (mean_demand + 1e-8)

        # Event sensitivity: avg demand on event days vs. non-event days
        event_mask = group["event_name"] != ""
        event_mean = y[event_mask.values].mean() if event_mask.any() else mean_demand
        non_event_mean = y[~event_mask.values].mean() if (~event_mask).any() else mean_demand
        event_sensitivity = event_mean / (non_event_mean + 1e-8)

        # Trend slope (linear regression coefficient, normalised)
        if len(t) > 1:
            slope = np.polyfit(t, y, 1)[0]
            trend_slope = slope / (mean_demand + 1e-8)
        else:
            trend_slope = 0.0

        # Seasonality strength: std of monthly means / overall std
        monthly_means = group.groupby(group["ds"].dt.month)["y"].mean()
        seasonality = monthly_means.std() / (y.std() + 1e-8)

        features[uid] = {
            "mean_demand":       mean_demand,
            "cv":                cv,
            "event_sensitivity": event_sensitivity,
            "trend_slope":       trend_slope,
            "seasonality":       seasonality,
            "category":          group["category"].iloc[0],
            "store":             group["store"].iloc[0],
        }

    return pd.DataFrame.from_dict(features, orient="index")


def cluster_products(
    feature_df: pd.DataFrame,
    min_cluster_size: int = 3,
    min_samples: int = 2,
) -> pd.DataFrame:
    """
    Run HDBSCAN on the numeric feature columns.
    Returns feature_df with added columns: cluster_id, segment_label.
    """
    numeric_cols = ["mean_demand", "cv", "event_sensitivity", "trend_slope", "seasonality"]
    X = feature_df[numeric_cols].values

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    clusterer = hdbscan.HDBSCAN(
        min_cluster_size=min_cluster_size,
        min_samples=min_samples,
        metric="euclidean",
        cluster_selection_method="eom",
    )
    labels = clusterer.fit_predict(X_scaled)

    result = feature_df.copy()
    result["cluster_id"] = labels
    result["segment_label"] = _infer_labels(result, labels)
    result["cluster_prob"] = clusterer.probabilities_

    return result


def _infer_labels(df: pd.DataFrame, labels: np.ndarray) -> pd.Series:
    """
    Map numeric cluster IDs to human-readable segment names
    by inspecting each cluster's centroid profile.
    """
    label_series = pd.Series("", index=df.index)

    for cluster_id in set(labels):
        if cluster_id == -1:
            label_series[df["cluster_id"] == cluster_id] = SEGMENT_LABELS["noise"]
            continue

        mask = df["cluster_id"] == cluster_id
        centroid = df.loc[mask, ["mean_demand", "cv", "trend_slope", "event_sensitivity"]].mean()

        high_volume  = centroid["mean_demand"] > df["mean_demand"].median()
        high_cv      = centroid["cv"] > df["cv"].median()
        trending_up  = centroid["trend_slope"] > 0.01
        trending_dn  = centroid["trend_slope"] < -0.01
        event_driven = centroid["event_sensitivity"] > 1.5

        if trending_dn:
            label = SEGMENT_LABELS["declining"]
        elif trending_up and not high_cv:
            label = SEGMENT_LABELS["growing"]
        elif high_volume and event_driven:
            label = SEGMENT_LABELS["high_volume_volatile"]
        elif high_volume and not high_cv:
            label = SEGMENT_LABELS["high_volume_stable"]
        else:
            label = SEGMENT_LABELS["low_volume_volatile"]

        label_series[mask] = label

    return label_series


def get_cluster_profiles(clustered_df: pd.DataFrame) -> pd.DataFrame:
    """
    Summarise each segment: size, mean features, representative category.
    Used by the Streamlit tab for cluster profile cards.
    """
    numeric_cols = ["mean_demand", "cv", "event_sensitivity", "trend_slope", "seasonality"]
    profiles = (
        clustered_df.groupby("segment_label")
        .agg(
            count=("mean_demand", "count"),
            **{col: (col, "mean") for col in numeric_cols},
        )
        .round(3)
        .reset_index()
    )
    return profiles