"""Plot-level aggregation of per-tuber morphometrics.

Per (trial, Plot_std): total projected area and tuber count.
Per (trial, Plot_std, class): class-specific projected area, pivoted
wide to area_{B,A1,A2,OV}_in2 with NaN replaced by 0 for absent classes.
"""

from __future__ import annotations

import pandas as pd

CLASS_AREA_COLUMNS: tuple[str, ...] = (
    "area_B_in2",
    "area_A1_in2",
    "area_A2_in2",
    "area_OV_in2",
)


def aggregate_total_area(df: pd.DataFrame) -> pd.DataFrame:
    """One row per (trial, Plot_std) with total_area_in2 and n_tubers."""
    return (
        df.groupby(["trial", "Plot_std"], as_index=False)
        .agg(
            total_area_in2=("area_poly_in2", "sum"),
            n_tubers=("detection_id", "count"),
        )
    )


def aggregate_area_by_class_wide(df: pd.DataFrame) -> pd.DataFrame:
    """One row per (trial, Plot_std) with one column per USDA class area."""
    long = (
        df.groupby(["trial", "Plot_std", "class"], as_index=False, observed=True)
        .agg(class_area_in2=("area_poly_in2", "sum"))
    )
    wide = (
        long.pivot(index=["trial", "Plot_std"], columns="class", values="class_area_in2")
        .fillna(0)
        .reset_index()
        .rename(
            columns={
                "B": "area_B_in2",
                "A1": "area_A1_in2",
                "A2": "area_A2_in2",
                "OV": "area_OV_in2",
            }
        )
    )
    wide.columns.name = None
    for col in CLASS_AREA_COLUMNS:
        if col not in wide.columns:
            wide[col] = 0.0
    return wide


def aggregate_by_plot(df: pd.DataFrame) -> pd.DataFrame:
    """Combine total-area and class-area-wide aggregations on (trial, Plot_std)."""
    df_total = aggregate_total_area(df)
    df_wide = aggregate_area_by_class_wide(df)
    return df_total.merge(df_wide, on=["trial", "Plot_std"], how="left")