import numpy as np
from sklearn.model_selection import StratifiedKFold, StratifiedGroupKFold, StratifiedShuffleSplit
import warnings

# Function to suppress the traceback of warnings
def custom_formatwarning(msg, *args, **kwargs):
    return str(msg) + '\n'

warnings.formatwarning = custom_formatwarning
warnings.filterwarnings('default', category=UserWarning)

# StratifiedKFold example
print("StratifiedKFold Example:")
X = np.random.rand(10, 2)  # 10 samples, 2 features
y_skf = np.zeros(10)  # Target variable with one class
skf = StratifiedKFold(n_splits=3)
for train_index, test_index in skf.split(X, y_skf):
    print(f"TRAIN indices: {train_index}, TEST indices: {test_index}")

# StratifiedGroupKFold example
print("\nStratifiedGroupKFold Example:")
X = np.random.rand(15, 2)  # 15 samples, 2 features
y_sgkf = np.ones(15)  # Target variable with one class
groups = np.array([1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7])
sgkf = StratifiedGroupKFold(n_splits=3)
for train_index, test_index in sgkf.split(X, y_sgkf, groups=groups):
    print(f"TRAIN indices: {train_index}, TEST indices: {test_index}")

# StratifiedShuffleSplit example
print("\nStratifiedShuffleSplit Example:")
X = np.random.rand(20, 2)  # 20 samples, 2 features
y_sss = np.zeros(20)  # Target variable with one class
sss = StratifiedShuffleSplit(n_splits=5, test_size=0.5, random_state=42)
for train_index, test_index in sss.split(X, y_sss):
    print(f"TRAIN indices: {train_index}, TEST indices: {test_index}")

Only one class present in y. StratifiedKFold is designed to be used with data that contains two or more classes. Consider using KFold instead.
Only one class present in y. StratifiedGroupKFold is designed to be used with data that contains two or more classes. The single-class scenario might not be suitable for stratified folds.
Only one class present in y. StratifiedShuffleSplit is designed to be used with data that contains two or more classes. Consider using ShuffleSplit instead.