from storai.schema.types import Metric


def compute_contribution(df, metric_col, dimension_col):
    grouped = df.groupby(dimension_col)[metric_col].sum()

    # Skip if only 1 category
    if grouped.nunique() <= 1:
        return []

    total = grouped.sum()  # Total contribution across all categories
    if total == 0:
        return []

    top_entity = grouped.idxmax()  # Get the category with the highest contribution
    share = grouped.max() / total  # Calculate the share of the top category

    return [
        Metric(
            type="contribution",
            name=metric_col,
            value=float(share),
            context={
                "dimension": dimension_col,
                "entity": top_entity,
                "total": float(total),
                "num_groups": len(grouped),
            },
        )
    ]
