from __future__ import annotations

import pandas as pd

from macro_lactone_toolkit import MacroLactoneAnalyzer


def stratified_sample_by_ring_size(
    df: pd.DataFrame,
    sample_ratio: float,
    smiles_col: str = "smiles",
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Perform stratified sampling by ring size.

    First classifies all molecules, then samples 10% from each ring size layer.
    """
    analyzer = MacroLactoneAnalyzer()

    # Classify all molecules
    classifications = []
    ring_sizes = []

    for smiles in df[smiles_col]:
        result = analyzer.classify_macrocycle(smiles)
        classifications.append(result.classification)
        ring_sizes.append(result.ring_size)

    df = df.copy()
    df["_classification"] = classifications
    df["_ring_size"] = ring_sizes

    # Group by ring size and sample from each group
    sampled_groups = []

    for ring_size in range(12, 21):
        group = df[df["_ring_size"] == ring_size]
        if len(group) > 0:
            n_samples = max(1, int(len(group) * sample_ratio))
            sampled = group.sample(n=min(n_samples, len(group)), random_state=random_state)
            sampled_groups.append(sampled)

    # Also sample from unknown ring size (None)
    unknown_group = df[df["_ring_size"].isna()]
    if len(unknown_group) > 0:
        n_samples = max(1, int(len(unknown_group) * sample_ratio))
        sampled = unknown_group.sample(n=min(n_samples, len(unknown_group)), random_state=random_state)
        sampled_groups.append(sampled)

    if not sampled_groups:
        return pd.DataFrame()

    result = pd.concat(sampled_groups, ignore_index=True)
    return result
