API Reference¶

@pd.api.extensions.register_dataframe_accessor("epi")
class EpiAccessor:
    """
    Pandas accessor for epidemiological analysis on isolate datasets.

    Provides methods for generating epidemic curves, calculating prevalence,
    diversity, and incidence, and identifying transmission clusters.
    """
    def __init__(self, pandas_obj: pd.DataFrame):
        self._obj = pandas_obj

    # --- Shiny UI State Checkers ---

    @property
    def has_temporal(self) -> bool:
        """Checks if the dataset contains valid temporal data (any 'temporal_' column)."""
        cols = self._obj.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").columns
        return bool(len(cols) > 0 and self._obj[cols[0]].notna().any())

    @property
    def has_spatial(self) -> bool:
        """Checks if the dataset contains valid spatial coordinates (latitude and longitude)."""
        return 'latitude' in self._obj.columns and 'longitude' in self._obj.columns

    # --- Spatiotemporal Helpers ---

    @property
    def temporal(self) -> pd.DataFrame:
        """
        Returns a DataFrame of all temporal columns, with the prefix removed.
        """
        return self._obj.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").rename(columns=lambda c: c.replace(f'{Domain.TEMPORAL.value}_', '', 1))

    @property
    def temporal_resolution(self) -> pd.DataFrame:
        """
        Returns a DataFrame of all temporal resolution columns.
        """
        return self._obj.filter(regex=f"^{Domain.TEMPORAL_RES.value}_").rename(columns=lambda c: c.replace(f'{Domain.TEMPORAL_RES.value}_', '', 1))

    @property
    def spatial(self) -> pd.DataFrame:
        """
        Returns the core spatial coordinates (latitude, longitude).

        Raises:
            ValueError: If spatial columns are missing.
        """
        if not self.has_spatial:
            raise ValueError("Spatial columns ('latitude', 'longitude') are missing.")
        return self._obj[['latitude', 'longitude']].astype("Float64")

    # --- Time Series / Epidemic Curve Methods ---

    def _get_spatiotemporal_arrays(self, temporal_col: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Helper to extract and format coordinates and dates for spatial clustering."""
        df = self._obj

        if not pd.api.types.is_datetime64_any_dtype(df[temporal_col]):
            raise TypeError(f"Temporal column '{temporal_col}' must be a datetime type. Ensure data is parsed via seroepi.io.")

        # Schema guarantees datetime64, so we just safely strip timezones if present
        date_series = df[temporal_col].dt.tz_localize(None)
        # Coerce coordinates to standard numpy floats
        coords = np.radians(df[['latitude', 'longitude']].astype(float).values)

        # Convert dates to raw days (use float to allow NaNs for missing dates)
        raw_dates = np.full(len(df), np.nan)
        date_mask = date_series.notna().values
        raw_dates[date_mask] = date_series[date_mask].values.astype('datetime64[D]').astype(float)

        # Create a boolean mask of rows that have all required spatiotemporal data
        valid_mask = ~(np.isnan(coords[:, 0]) | np.isnan(coords[:, 1]) | np.isnan(raw_dates))
        return coords, raw_dates, valid_mask

    def epidemic_curve(self, freq: Union[str, TemporalResolution] = TemporalResolution.WEEK,
                       stratify_by: str = None, temporal_col: str = None) -> pd.DataFrame:
        """
        Generates a time-series DataFrame for plotting epidemic curves.

        Args:
            freq: Time frequency for resampling (e.g., TimeResolution.MONTH, 'ME', 'YE').
                Defaults to TimeResolution.WEEK.
            stratify_by: Column name to group by before resampling.

        Returns:
            A DataFrame with counts of isolates per time interval.

        Raises:
            ValueError: If temporal data is not available.
        """
        if not self.has_temporal:
            raise ValueError("Cannot generate epi curve: No temporal data available.")

        df = self._obj.copy()

        if isinstance(freq, TemporalResolution):
            freq_val = freq.pandas_offset
        else:
            freq_val = freq

        if temporal_col is None:
            temporal_col = df.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").columns[0]
        elif not temporal_col.startswith(f"{Domain.TEMPORAL.value}_"):
            temporal_col = f"{Domain.TEMPORAL.value}_{temporal_col}"

        if not pd.api.types.is_datetime64_any_dtype(df[temporal_col]):
            raise TypeError(f"Temporal column '{temporal_col}' must be a datetime type. Ensure data is parsed via seroepi.io.")

        # Set time as index for resampling
        df = df.set_index(temporal_col)

        if stratify_by:
            curve = df.groupby(stratify_by, observed=True).resample(freq_val).size().unstack(level=0, fill_value=0)
        else:
            curve = df.resample(freq_val).size().to_frame(name='count')

        return curve.reset_index()

    @property
    def metadata_columns(self) -> list[str]:
        """Returns the raw names of user-uploaded clinical/metadata columns."""
        return self._obj.filter(regex="^meta_").columns.tolist()

    @property
    def ui_metadata_columns(self) -> list[str]:
        """Returns clean metadata names (without 'meta_' prefix) for UI display."""
        return [c.replace('meta_', '', 1) for c in self.metadata_columns]

    @property
    def genotypes(self) -> list[str]:
        """Compiles a list of all genetic variables (Core + Accessory traits)."""
        # Grab all the dynamically prefixed traits (genotype, phenotype, amr, virulence)
        return self._obj.filter(regex=f"^({Domain.GENOTYPE.value}|{Domain.PHENOTYPE.value}|{Domain.AMR.value}|{Domain.VIRULENCE.value})_").columns.tolist()

    @property
    def stratify_cols(self) -> list[str]:
        """Returns columns suitable for stratification (excluding QC, metadata, and high-cardinality/internal cols)."""
        exclude_strat = ['sample_id', 'latitude', 'longitude']
        return [c for c in self._obj.columns if not c.startswith((f'{Domain.QC.value}_', 'meta_', f'{Domain.SPATIAL_RES.value}_', f'{Domain.TEMPORAL_RES.value}_')) and c not in exclude_strat]

    @property
    def cluster_cols(self) -> list[str]:
        """Returns columns suitable for cluster adjustment (e.g., transmission clusters and ST)."""
        return [c for c in self._obj.columns if c.startswith(f"{Domain.CLUSTER.value}_") or c.endswith('_ST')]

    @staticmethod
    def _calculate_events_and_denoms(
        df: pd.DataFrame,
        trait_col: str,
        denom_cols: list[str],
        trait_strata: list[str],
        cluster_col: str,
        negative_indicator: Union[str, list[str]],
        event_name: str,
        denom_name: str
    ) -> tuple[Union[pd.Series, int], Union[pd.Series, int]]:
        """Helper to calculate numerators and denominators for aggregated metrics."""
        if trait_col:
            valid_df = df.dropna(subset=[trait_col]).copy()

            if pd.api.types.is_bool_dtype(valid_df[trait_col]):
                valid_df['_trait_bool'] = valid_df[trait_col]
            else:
                neg_list = [negative_indicator] if isinstance(negative_indicator, str) else negative_indicator
                valid_df['_trait_bool'] = ~valid_df[trait_col].isin(neg_list)

            if cluster_col:
                denoms = valid_df.groupby(denom_cols, observed=True)[cluster_col].nunique().rename(denom_name) if denom_cols else valid_df[cluster_col].nunique()
                events = valid_df[valid_df['_trait_bool']].groupby(trait_strata, observed=True)[cluster_col].nunique().rename(event_name) if trait_strata else valid_df[valid_df['_trait_bool']][cluster_col].nunique()
            else:
                denoms = valid_df.groupby(denom_cols, observed=True).size().rename(denom_name) if denom_cols else len(valid_df)
                events = valid_df.groupby(trait_strata, observed=True)['_trait_bool'].sum().rename(event_name) if trait_strata else valid_df['_trait_bool'].sum()

        else:
            valid_df = df.dropna(subset=[trait_strata[-1]]).copy()

            if cluster_col:
                denoms = valid_df.groupby(denom_cols, observed=True)[cluster_col].nunique().rename(denom_name) if denom_cols else valid_df[cluster_col].nunique()
                events = valid_df.groupby(trait_strata, observed=True)[cluster_col].nunique().rename(event_name)
            else:
                denoms = valid_df.groupby(denom_cols, observed=True).size().rename(denom_name) if denom_cols else len(valid_df)
                events = valid_df.groupby(trait_strata, observed=True).size().rename(event_name)

        return events, denoms

    def aggregate_prevalence(self, stratify_by: list[str], trait_col: str = None,
                             cluster_col: str = None, negative_indicator: Union[str, list[str]] = '-',
                             pad_zeros: bool = False) -> pd.DataFrame:
        """
        Aggregates data to calculate event counts and denominators for prevalence.

        Supports both trait prevalence (presence/absence of a marker) and
        compositional prevalence (distribution of variants within a locus).

        Args:
            stratify_by: Columns to group by (e.g., ['spatial']).
            trait_col: The column containing the trait/marker to measure.
                If None, compositional prevalence is calculated for the last
                column in `stratify_by`.
            cluster_col: Column containing cluster IDs to adjust for (e.g., nosocomial outbreaks).
            negative_indicator: Value(s) representing the absence of a trait.
                Defaults to '-'.
            pad_zeros: If True, pads missing combinations of strata with zero counts.
                Essential for spatial/hierarchical models. If False, only includes
                observed combinations for efficiency. Defaults to False.

        Returns:
            An aggregated DataFrame with 'event' and 'n' columns.

        Raises:
            ValueError: If compositional prevalence is requested without enough strata.
        """
        df = self._obj.copy()

        if trait_col:
            denom_cols = stratify_by
            trait_strata = stratify_by
        else:
            if len(stratify_by) < 1:
                raise ValueError("Compositional prevalence requires at least 1 stratify_by column.")
            denom_cols = stratify_by[:-1]
            trait_strata = stratify_by

        events, denoms = self._calculate_events_and_denoms(
            df=df,
            trait_col=trait_col,
            denom_cols=denom_cols,
            trait_strata=trait_strata,
            cluster_col=cluster_col,
            negative_indicator=negative_indicator,
            event_name='event',
            denom_name='n'
        )

        # Expand Grid
        if len(stratify_by) > 0:
            if pad_zeros:
                # Full Cartesian expansion (padding zeroes)
                unique_levels = [df[col].dropna().unique() for col in stratify_by]
                base_index = pd.MultiIndex.from_product(unique_levels, names=stratify_by)
            else:
                # Efficiently use only the observed strata combinations
                base_index = denoms.index if trait_col else events.index

            agg_df = pd.DataFrame(index=base_index).join(events, how='left').fillna({'event': 0}).reset_index()

            # Safely map denominators
            if len(denom_cols) == 0:
                agg_df['n'] = denoms
            else:
                agg_df = agg_df.set_index(denom_cols)
                agg_df['n'] = denoms
                agg_df = agg_df.reset_index().fillna({'n': 0})

            if not pad_zeros:
                agg_df = agg_df[agg_df['n'] > 0].copy()
        else:
            agg_df = pd.DataFrame({'event': [events], 'n': [denoms]})

        # Standardize to 'target' for uniform downstream handling
        if trait_col:
            agg_df['target'] = trait_col
        else:
            agg_df = agg_df.rename(columns={stratify_by[-1]: 'target'})

        agg_df.attrs = self._obj.attrs.copy()
        agg_df.attrs['metric_meta'] = {
            "metric_type": MetricType.PREVALENCE,
            "stratified_by": denom_cols,
            "trait": trait_col if trait_col else stratify_by[-1],
            "aggregation_type": AggregationType.TRAIT if trait_col else AggregationType.COMPOSITIONAL,
            "adjusted_for": cluster_col,
            "is_zero_padded": pad_zeros
        }

        return agg_df

    def aggregate_diversity(self, stratify_by: list[str], trait_col: str = None,
                            cluster_col: str = None, negative_indicator: Union[str, list[str]] = '-',
                            pad_zeros: bool = False) -> pd.DataFrame:
        """
        Aggregates data to calculate counts for diversity analysis (e.g., Shannon index).

        Args:
            stratify_by: Columns to group by.
            trait_col: The locus or trait to measure diversity for.
            cluster_col: Optional cluster column to adjust for.
            negative_indicator: Value(s) to exclude from diversity counts.
            pad_zeros: If True, pads missing combinations of strata with zero counts.
                Defaults to False.

        Returns:
            A DataFrame with 'variant_count' and 'n_total'.

        Raises:
            ValueError: If compositional diversity is requested without enough strata.
        """
        df = self._obj.copy()

        is_trait = True if trait_col else False
        if trait_col:
            groupers = stratify_by
            trait_strata = stratify_by + [trait_col]
            valid_df = df.dropna(subset=[trait_col]).copy()

            # For trait diversity, we often want to strip out the "absence" indicators
            # so they don't count as a diversity variant mathematically
            if not pd.api.types.is_bool_dtype(valid_df[trait_col]):
                neg_list = [negative_indicator] if isinstance(negative_indicator, str) else negative_indicator
                valid_df = valid_df[~valid_df[trait_col].isin(neg_list)]
        else:
            if not stratify_by:
                raise ValueError("Compositional diversity requires at least 1 stratify_by column.")
            groupers = stratify_by[:-1]
            trait_col = stratify_by[-1]
            trait_strata = stratify_by
            valid_df = df.dropna(subset=[trait_col]).copy()

        if cluster_col:
            div_df = valid_df.groupby(trait_strata, observed=True)[cluster_col].nunique().reset_index()
            div_df = div_df.rename(columns={cluster_col: 'variant_count'})
        else:
            div_df = valid_df.groupby(trait_strata, observed=True).size().reset_index(name='variant_count')

        if pad_zeros and trait_strata:
            unique_levels = [df[col].dropna().unique() for col in trait_strata]
            expanded_index = pd.MultiIndex.from_product(unique_levels, names=trait_strata)
            div_df = div_df.set_index(trait_strata).reindex(expanded_index, fill_value=0).reset_index()

        if groupers:
            div_df['n_total'] = div_df.groupby(groupers, observed=True)['variant_count'].transform('sum')
        else:
            div_df['n_total'] = div_df['variant_count'].sum()

        # Standardize to 'target'
        div_df = div_df.rename(columns={trait_col: 'target'})
        if is_trait:
            div_df['target'] = trait_col

        div_df.attrs = self._obj.attrs.copy()
        div_df.attrs["metric_meta"] = {
            "metric_type": MetricType.DIVERSITY,
            "stratified_by": groupers,
            "trait": trait_col,
            "aggregation_type": AggregationType.TRAIT if trait_col else AggregationType.COMPOSITIONAL,
            "adjusted_for": cluster_col,
            "is_zero_padded": pad_zeros
        }

        return div_df

    def aggregate_incidence(self, stratify_by: list[str], trait_col: str = None, freq: Union[str, TemporalResolution] = TemporalResolution.MONTH,
                            cluster_col: str = None, negative_indicator: Union[str, list[str]] = '-',
                            pad_zeros: bool = False, temporal_col: str = None) -> pd.DataFrame:
        """
        Aggregates data for time-series incidence analysis.

        Args:
            stratify_by: Columns to group by.
            trait_col: The marker to measure incidence for.
            freq: Time frequency for binning (e.g., TimeResolution.MONTH, 'ME'). Defaults to TimeResolution.MONTH.
            cluster_col: Optional cluster column to adjust for.
            negative_indicator: Value(s) representing absence.
            pad_zeros: If True, pads missing combinations of strata. If False,
                maintains unbroken time grids only for observed strata combinations.

        Returns:
            A DataFrame with 'variant_count', 'total_sequenced', and binned dates.

        Raises:
            ValueError: If temporal data is missing or inappropriate strata provided.
        """
        df = self._obj.copy()

        if temporal_col is None:
            temporal_cols = df.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").columns
            if not len(temporal_cols):
                raise ValueError("Incidence aggregation requires a valid temporal column.")
            temporal_col = temporal_cols[0]
        elif not temporal_col.startswith(f"{Domain.TEMPORAL.value}_"):
            temporal_col = f"{Domain.TEMPORAL.value}_{temporal_col}"

        if temporal_col not in df.columns or not pd.api.types.is_datetime64_any_dtype(df[temporal_col]):
            raise ValueError(f"Incidence aggregation requires a valid datetime64 temporal column. '{temporal_col}' invalid.")

        # Translate Pandas 2.2+ point offsets ('ME') to period spans ('M')
        if isinstance(freq, TemporalResolution):
            period_freq = freq.pandas_period
            stored_freq = freq.value
        else:
            period_freq = freq.replace('ME', 'M').replace('YE', 'Y')
            stored_freq = freq

        # Snap dates to the requested frequency bin using the safe string
        df['date_bin'] = df[temporal_col].dt.to_period(period_freq).dt.to_timestamp()

        if trait_col:
            denom_cols = ['date_bin'] + stratify_by
            trait_strata = ['date_bin'] + stratify_by
        else:
            if not stratify_by:
                raise ValueError("Compositional incidence requires at least 1 stratify_by column.")

            denom_cols = ['date_bin'] + stratify_by[:-1]
            trait_strata = ['date_bin'] + stratify_by

        events, denoms = self._calculate_events_and_denoms(
            df=df,
            trait_col=trait_col,
            denom_cols=denom_cols,
            trait_strata=trait_strata,
            cluster_col=cluster_col,
            negative_indicator=negative_indicator,
            event_name='variant_count',
            denom_name='total_sequenced'
        )

        # --- TIME GRID EXPANSION ---
        # Generate an unbroken sequence of dates from the earliest to the latest record
        min_date, max_date = df['date_bin'].min(), df['date_bin'].max()
        all_dates = pd.period_range(min_date, max_date, freq=period_freq).to_timestamp().tolist() if not pd.isna(
            min_date) else []

        if pad_zeros:
            unique_levels = [all_dates] + [df[col].dropna().unique() for col in trait_strata[1:]]
            expanded_index = pd.MultiIndex.from_product(unique_levels, names=trait_strata)
        else:
            if len(trait_strata) > 1:
                # Only use strata combinations that actually appear in the data
                observed_strata = df[trait_strata[1:]].dropna().drop_duplicates()
                dates_df = pd.DataFrame({trait_strata[0]: all_dates})
                expanded_df = dates_df.merge(observed_strata, how='cross')
                expanded_index = pd.MultiIndex.from_frame(expanded_df[trait_strata])
            else:
                expanded_index = pd.Index(all_dates, name=trait_strata[0])

        inc_df = pd.DataFrame(index=expanded_index).join(events, how='left').fillna({'variant_count': 0}).reset_index()

        inc_df = inc_df.set_index(denom_cols)
        inc_df['total_sequenced'] = denoms
        inc_df = inc_df.reset_index().fillna({'total_sequenced': 0})

        # Unlike Prevalence, we do NOT drop rows where total_sequenced == 0.
        # A true 0 sequence volume is critical information for an epicurve gap.
        inc_df = inc_df.rename(columns={'date_bin': 'date'})

        # Standardize to 'target'
        if trait_col:
            inc_df['target'] = trait_col
        else:
            inc_df = inc_df.rename(columns={stratify_by[-1]: 'target'})

        inc_df.attrs = self._obj.attrs.copy()
        inc_df.attrs['metric_meta'] = {
            "metric_type": MetricType.INCIDENCE,
            "stratified_by": stratify_by if trait_col else stratify_by[:-1],
            "trait": trait_col if trait_col else stratify_by[-1],
            "aggregation_type": AggregationType.TRAIT if trait_col else AggregationType.COMPOSITIONAL,
            "freq": stored_freq,
            "adjusted_for": cluster_col,
            "is_zero_padded": pad_zeros
        }

        return inc_df

    def transmission_network(
            self,
            clone_col: str,
            spatial_threshold_km: float = 10.0,
            temporal_threshold_days: int = 20,
            temporal_col: str = None
    ) -> TransmissionDistances:
        """
        Builds a sparse adjacency graph of transmission links.

        Args:
            clone_col: Column containing clone IDs (e.g., 'ST' or a custom cluster).
            spatial_threshold_km: Maximum distance in kilometers. Defaults to 10.0.
            temporal_threshold_days: Maximum time difference in days. Defaults to 20.

        Returns:
            A TransmissionDistances object representing the outbreak network.

        Raises:
            KeyError: If required columns ('latitude', 'longitude', 'date') are missing.
        """
        df = self._obj

        # 1. Validation Checks
        if clone_col not in df.columns:
            raise KeyError(f"Clone column '{clone_col}' not found in DataFrame.")

        # Intelligently check for your geo accessor/columns
        if 'latitude' not in df.columns or 'longitude' not in df.columns:
            raise KeyError("Spatial clustering requires 'latitude' and 'longitude' columns. "
                           "Ensure geo accessors have parsed coordinates.")

        if temporal_col is None:
            temporal_cols = df.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").columns
            if not len(temporal_cols):
                raise KeyError("A temporal column is required for temporal clustering.")
            temporal_col = temporal_cols[0]
        elif not temporal_col.startswith(f"{Domain.TEMPORAL.value}_"):
            temporal_col = f"{Domain.TEMPORAL.value}_{temporal_col}"

        if temporal_col not in df.columns:
            raise KeyError(f"Temporal column '{temporal_col}' not found.")

        df = self._obj
        coords, raw_dates, _ = self._get_spatiotemporal_arrays(temporal_col)

        return TransmissionDistances.from_spatiotemporal(
            sample_ids=df['sample_id'],
            coords=coords,
            dates=raw_dates,
            clones=df[clone_col].values,
            spatial_threshold_km=spatial_threshold_km,
            temporal_threshold_days=temporal_threshold_days
        )

    def transmission_clusters(
            self,
            clone_col: str,
            spatial_threshold_km: float = 10.0,
            temporal_threshold_days: int = 20,
            temporal_col: str = None,
            network: TransmissionDistances = None
    ) -> pd.Series:
        """Extracts and formats categorical cluster labels from the transmission network."""
        df = self._obj

        if temporal_col is None:
            temporal_cols = df.filter(regex=f"^{Domain.TEMPORAL.value}_(?!res_)").columns
            if not len(temporal_cols):
                raise KeyError("A temporal column is required for temporal clustering.")
            temporal_col = temporal_cols[0]
        elif not temporal_col.startswith(f"{Domain.TEMPORAL.value}_"):
            temporal_col = f"{Domain.TEMPORAL.value}_{temporal_col}"

        if network is None:
            network = self.transmission_network(clone_col, spatial_threshold_km, temporal_threshold_days, temporal_col)

        labels = network.get_clusters()

        _, _, valid_mask = self._get_spatiotemporal_arrays(temporal_col)
        clone_mask = df[clone_col].notna().values

        labels_array = labels.astype(float).to_numpy(copy=True)
        labels_array[~valid_mask] = np.nan
        labels_array[~clone_mask] = np.nan

        res = pd.Series(labels_array, index=df.index, dtype="Int64",
                        name=f'{Domain.CLUSTER.value}_transmission_{spatial_threshold_km}km_{temporal_threshold_days}days')
        return res.astype("category").cat.as_ordered()