import pandas as pd

class Sampler:
    def __init__(self, sampling_config: dict):
        self.sampling_config = sampling_config

    def sample(self, df: pd.DataFrame) -> pd.DataFrame:
        strategy = self.sampling_config.get("type")

        if strategy == "random":
            return self.random_sample(df)
        elif strategy == "stratified":
            return self.stratified_sample(df)
        elif strategy == "stratified-proportional":
            return self.stratified_proportional_sample(df)
        else:
            raise ValueError(f"Unsupported sampling strategy: {strategy}")

    def random_sample(self, df: pd.DataFrame) -> pd.DataFrame:
        sample_size = self.sampling_config.get("sample_size")
        if sample_size is None:
            raise ValueError("`sample_size` must be specified for random sampling.")
        return df.sample(n=min(sample_size, len(df)), random_state=42)

    def stratified_sample(self, df: pd.DataFrame) -> pd.DataFrame:
        by = self.sampling_config.get("by")
        ratio = self.sampling_config.get("ratio")

        if not by:
            raise ValueError("`by` must be specified for stratified sampling.")

        grouped = df.groupby(by)
        sampled_groups = []

        # Auto-compute equal ratio if ratio not passed
        if ratio is None:
            default_ratio = self.sampling_config.get("default_per_group", 10)
            print(f"[INFO] No ratio provided. Using default ratio of {default_ratio} per group.")
            for group_key, group_df in grouped:
                n = min(default_ratio, len(group_df))
                if n > 0:
                    sampled_groups.append(group_df.sample(n=n, random_state=42))
        else:
            for group_key, group_df in grouped:
                lookup_key = ":".join(map(str, group_key)) if isinstance(group_key, tuple) else str(group_key)
                group_n = ratio.get(lookup_key, 0) if isinstance(ratio, dict) else ratio
                n = min(group_n, len(group_df))
                if n > 0:
                    sampled_groups.append(group_df.sample(n=n, random_state=42))
                    print(f"[STRATIFIED] {lookup_key} → requested: {group_n}, sampled: {n}")

        return pd.concat(sampled_groups, ignore_index=True) if sampled_groups else pd.DataFrame(columns=df.columns)

    def stratified_proportional_sample(self, df: pd.DataFrame) -> pd.DataFrame:
        by = self.sampling_config.get("by")
        proportion = self.sampling_config.get("proportion")

        if not by:
            raise ValueError("`by` must be specified for stratified-proportional sampling.")

        grouped = df.groupby(by)
        sampled_groups = []

        # Auto-equal proportions if not passed
        if proportion is None:
            equal_p = 1 / len(grouped)
            print(f"[INFO] No proportion provided. Using equal {round(equal_p * 100)}% per group.")
            for group_key, group_df in grouped:
                n = int(len(group_df) * equal_p)
                if n > 0:
                    sampled_groups.append(group_df.sample(n=n, random_state=42))
        else:
            for group_key, group_df in grouped:
                lookup_key = ":".join(map(str, group_key)) if isinstance(group_key, tuple) else str(group_key)
                p = proportion.get(lookup_key, 0) if isinstance(proportion, dict) else proportion
                n = int(len(group_df) * p)
                if n > 0:
                    sampled_groups.append(group_df.sample(n=n, random_state=42))
                    print(f"[STRAT-PROP] {lookup_key} → {p*100}% → sampled: {n}")

        return pd.concat(sampled_groups, ignore_index=True) if sampled_groups else pd.DataFrame(columns=df.columns)