"""USDA size-class assignment from the minimum in-plane tuber dimension.

Reference: https://www.ams.usda.gov/grades-standards/potatoes-grades-and-standards
Classification is based on min(width_in, height_in) of the polygon
bounding box (the inMinDim column). The PO ("pickout") class is not
assigned here, since PO is a quality grade rather than a size class.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

# Right-open bins: [B] < 1.875 ≤ [A1] < 2.5625 ≤ [A2] < 3.25 ≤ [OV]
USDA_BIN_EDGES_IN = [-np.inf, 1.875, 2.5625, 3.25, np.inf]
USDA_LABELS = ["B", "A1", "A2", "OV"]


def assign_usda_class(df: pd.DataFrame, dim_col: str = "inMinDim") -> pd.DataFrame:
    """Add a categorical `class` column with USDA size label per tuber."""
    df = df.copy()
    df["class"] = pd.cut(
        df[dim_col],
        bins=USDA_BIN_EDGES_IN,
        labels=USDA_LABELS,
        right=False,
    )
    return df