import pandas as pd
import numpy as np


def performance_based_split(
    subject_ids,
    performance_df,
    split_ratio=0.33,
    threshold=0.01,
    max_iter=100,
    random_seed=None
):
    """
    Split subjects into two groups based on performance scores with balanced means.
    
    Parameters
    ----------
    subject_ids : array-like
        List or array of subject IDs present in your dataset
    performance_df : pd.DataFrame
        DataFrame containing 'subjectID' and 'overall_score' columns
    split_ratio : float, default=0.33
        Proportion of subjects for the smaller group (0 < split_ratio < 1)
    threshold : float, default=0.01
        Target difference threshold between group means
    max_iter : int, default=100
        Maximum number of swap iterations
    random_seed : int, optional
        Random seed for reproducibility
        
    Returns
    -------
    group_small_ids : np.ndarray
        Subject IDs for the smaller group
    group_large_ids : np.ndarray
        Subject IDs for the larger group
    score_diff : float
        Final absolute difference between group means
        
    Raises
    ------
    ValueError
        If subjects are missing performance scores or no subjects match
    """
    if random_seed is not None:
        np.random.seed(random_seed)
    
    # Filter performance data
    perf_filtered = performance_df[
        performance_df["subjectID"].isin(subject_ids)
    ][["subjectID", "overall_score"]]
    
    # Merge to get only subjects present in both dataset and performance file
    merged = (
        pd.DataFrame({"subjectID": subject_ids})
        .merge(perf_filtered, on="subjectID", how="inner")
    )
    
    if len(merged) == 0:
        raise ValueError("No subjects found in both dataset and performance file.")
    
    # Check for missing scores
    if merged["overall_score"].isna().any():
        raise ValueError("Missing score values for some subjects.")
    
    merged_sorted = merged.sort_values("overall_score", ascending=False).reset_index(drop=True)
    
    scores = merged_sorted["overall_score"].values
    n_total = len(merged_sorted)
    n_small = int(n_total * split_ratio)
    n_large = n_total - n_small
    
    # Initial random split
    idx = np.arange(n_total)
    np.random.shuffle(idx)
    
    small_idx = idx[:n_small]
    large_idx = idx[n_small:]
    
    def score_diff(small_idx, large_idx):
        return abs(scores[small_idx].mean() - scores[large_idx].mean())
    
    diff = score_diff(small_idx, large_idx)
    count = 0
    
    # Optimize via random swaps
    while diff > threshold and count < max_iter:
        si = np.random.choice(small_idx)
        li = np.random.choice(large_idx)
        
        new_small_idx = small_idx.copy()
        new_large_idx = large_idx.copy()
        
        new_small_idx[new_small_idx == si] = li
        new_large_idx[new_large_idx == li] = si
        
        new_diff = score_diff(new_small_idx, new_large_idx)
        
        if new_diff < diff:
            small_idx = new_small_idx
            large_idx = new_large_idx
            diff = new_diff
        
        count += 1
    
    # Extract subject IDs
    group_small_ids = merged_sorted.loc[small_idx, "subjectID"].values
    group_large_ids = merged_sorted.loc[large_idx, "subjectID"].values
    
    return group_small_ids, group_large_ids, diff


# Example usage for 2-way split (train/test)
# subjects = df["subjectID"].unique()
# performance_df = pd.read_csv("performance.csv")
# 
# train_ids, test_ids, diff = performance_based_split(
#     subject_ids=subjects,
#     performance_df=performance_df,
#     split_ratio=0.2,
#     random_seed=42
# )

# Example usage for 3-way split (train/val/test)
# Step 1: Split into train and temp
# train_ids, temp_ids, diff1 = performance_based_split(
#     subject_ids=subjects,
#     performance_df=performance_df,
#     split_ratio=0.6,  # 60% train, 40% temp
#     random_seed=42
# )
# 
# Step 2: Split temp into val and test
# val_ids, test_ids, diff2 = performance_based_split(
#     subject_ids=temp_ids,
#     performance_df=performance_df,
#     split_ratio=0.5,  # 50/50 split of remaining 40%
#     random_seed=43
# )