import pandas as pd
from typing import Tuple, Dict, Union, List
from sklearn.utils import resample


def split(
    data: pd.DataFrame,
    test_size: float,
    seed: int = 42
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Splits data into train and test sets.

    Args:
        data: Input dataframe
        test_size: Fraction of data for test set (0.0 to 1.0)
        seed: Random seed for reproducibility

    Returns:
        Tuple of (train_df, test_df)
    """
    train_size = int(len(data) * (1 - test_size))
    shuffled = data.sample(frac=1, random_state=seed).reset_index(drop=True)
    return shuffled[:train_size], shuffled[train_size:]


def downsample(
    data: pd.DataFrame,
    sample_ratio: float,
    seed: int = 42
) -> pd.DataFrame:
    """
    Randomly downsample the dataset.

    Args:
       data: Input dataframe
       sample_ratio: Fraction to keep (0.0 to 1.0)
       seed: Random seed for reproducibility

    Returns:
       Downsampled dataframe
    """
    return data.sample(frac=sample_ratio, random_state=seed).reset_index(drop=True)


def rebalance(
    data: pd.DataFrame,
    target_distribution: Dict[str, Union[int, float]],
    seed: int = 42
) -> pd.DataFrame:
    """
    Rebalance dataset according to target distribution.

    Args:
        data: Input dataframe
        label_column: Name of the column containing labels
        target_distribution: Dict mapping label to target count (int) or fraction (float)
        seed: Random seed for reproducibility

    Returns:
        Rebalanced dataframe
    """
    import ast

    def parse_label(x):
        if isinstance(x, str):
            label_dict = ast.literal_eval(x)
        else:
            label_dict = x
        return next(k for k, v in label_dict.items() if v == 1)

    df = data.copy()
    label_column = '_temp_label'
    df[label_column] = df['completion'].apply(parse_label)

    # Determine rebalance mode based on the values
    values = list(target_distribution.values())
    if all(isinstance(v, int) for v in values):
        mode = "absolute"
    elif all(isinstance(v, (int, float)) for v in values):
        mode = "fraction"
    else:
        raise ValueError("All values must be int or float")

    grouped = df.groupby(label_column)
    balanced_dfs: List[pd.DataFrame] = []

    if mode == "absolute":
        for label, group in grouped:
            target_size = target_distribution[str(label)]
            balanced_dfs.append(_resample_group(group, target_size, seed))
    else:
        total_size = len(data)
        for label, group in grouped:
            fraction = target_distribution[str(label)]
            target_size = int(fraction * total_size)
            balanced_dfs.append(_resample_group(group, target_size, seed))

    result = pd.concat(balanced_dfs, ignore_index=True)
    result.sample(frac=1, random_state=seed).reset_index(drop=True)

    return result[['prompt', 'completion']]


def _resample_group(group: pd.DataFrame, target_size: int, seed: int) -> pd.DataFrame:
    """Helper function to resample a group to target size."""
    current_size = len(group)
    if current_size == target_size:
        return group

    replace = current_size < target_size
    return resample(group, n_samples=target_size, random_state=seed, replace=replace)
