import numpy as np
import pandas as pd
from sklearn import model_selection

def create_folds(data):
    # for small dataset, use Leave-One-Out Cross-Validation
    
    data["kfold"] = -1  # Initialize a new column called kfold with -1
    data = data.sample(frac=1).reset_index(drop=True)  # Shuffle the data
   
    # Implement Leave-One-Out Cross-Validation
    loo = model_selection.LeaveOneOut()
    for fold_idx, (train_idx, val_idx) in enumerate(loo.split(data)):
        data.loc[val_idx, "kfold"] = fold_idx

    return data

if __name__ == "__main__":
    df = pd.read_pickle("input/tidbcloud_dedicated_pricing_dataset.pkl")
    df = create_folds(df)
    print(f"Total number of folds: {df['kfold'].nunique()}")
    print(f"Total number of samples: {len(df)}")
    df.to_pickle("input/tidbcloud_dedicated_pricing_dataset.pkl")