"""QC filters applied to the per-tuber morphometrics dataframe.

Three steps:
- Drop ADC plot images whose ROI cropping removed substantial tuber area.
- Derive the PCA-based axis ratio from eccentricity.
- Filter out tubers with axis_ratio_pca above the pooled 99.5th percentile.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

# ADC plot images dropped because the cropping step removed substantial
# tuber area. Values are the raw "Plot" column (image_id[-8:]).
ADC_QC_DROP_PLOTS: tuple[str, ...] = (
    "5_ADC116",
    "5_ADC131",
    "5_ADC141",
    "5_ADC129",
    "5_ADC105",
    "5_ADC205",
)

AR_PCA_FILTER_QUANTILE: float = 0.995


def drop_qc_plots(
    df: pd.DataFrame,
    drop_plots: tuple[str, ...] = ADC_QC_DROP_PLOTS,
    plot_col: str = "Plot",
) -> pd.DataFrame:
    """Drop plots listed in drop_plots from the per-tuber dataframe."""
    return df[~df[plot_col].isin(drop_plots)].copy()


def derive_axis_ratio_pca(df: pd.DataFrame) -> pd.DataFrame:
    """Add axis_ratio_pca = 1 / sqrt(1 - eccentricity^2)."""
    df = df.copy()
    ecc = np.clip(df["eccentricity"], 0, 0.999999)
    df["axis_ratio_pca"] = 1.0 / np.sqrt(1.0 - ecc**2)
    return df


def filter_axis_ratio(
    df: pd.DataFrame,
    quantile: float = AR_PCA_FILTER_QUANTILE,
) -> pd.DataFrame:
    """Keep tubers with axis_ratio_pca at or below the pooled quantile threshold."""
    upper = df["axis_ratio_pca"].quantile(quantile)
    return df[df["axis_ratio_pca"] <= upper].copy()