# AUTOGENERATED! DO NOT EDIT! File to edit: nbks/06_create_folds.ipynb (unless otherwise specified).

__all__ = ['create_folds']

# Cell
from pathlib import Path
import pandas as pd
from sklearn.model_selection import StratifiedKFold
from fastcore.script import call_parse, Param
from fastcore.script import store_true, bool_arg

# Cell
@call_parse
def create_folds(
    path:        Param("Path to train CSV", Path),
    print_stats: Param("Print class distributions across folds?", store_true),
    save: Param("Save CSV with added folds", bool_arg)=True,
) -> pd.DataFrame:
    """Saves train CSV at `path` with 5-fold CV splits added.

    Optionally, print class distribution across folds.
    """
    # Load train CSV
    df = pd.read_csv(path)

    # Shuffle rows
    df = df.sample(frac=1.0, random_state=42).reset_index(drop=True)

    # Get class labels to stratify on
    lbls = df.apply(lambda r: df.columns[r==1].item(), axis=1)

    # Create 5 folds
    kf = StratifiedKFold(n_splits=5)
    for fold, (_, val_idxs) in enumerate(kf.split(df, lbls.values)):
        df.loc[val_idxs, "fold"] = fold

    # Correct dtype of folds column
    df.fold = df.fold.astype(int)

    if print_stats:
        stats_df = df.groupby("fold").describe()
        print("Proportion of each class out of total examples in each fold:")
        print(stats_df.iloc[:, stats_df.columns.get_level_values(1) == "mean"])

    if save:
        save_path = path.parent/"train_folds.csv"
        df.to_csv(save_path, index=False)
        print(f"Saved to {save_path}")

    return df