from pandas import read_csv
from sklearn.model_selection import KFold
from sklearn.preprocessing import LabelEncoder


def replace_values(df):
    columns = ["BI-RADS", "Margin", "Density", "Age"]
    for column in columns:
        df[column].fillna(value=df[column].mean(), inplace=True)
    return df


def process_na(df, action):
    if action == "drop":
        return df.dropna()
    elif action == "fill":
        return replace_values(df)
    else:
        print("Unknown action selected. The choices are: ")
        print("fill: fills the na values with the mean")
        print("drop: drops the na values")
        exit()


def encode_columns(df):
    label_encoder = LabelEncoder()
    encoded_df = df.copy()
    encoded_df["Shape"] = label_encoder.fit_transform(df["Shape"])
    encoded_df["Severity"] = label_encoder.fit_transform(df["Severity"])
    return encoded_df


def split_train_target(df):
    train_data = df.drop(columns=["Severity"])
    target_data = df["Severity"]
    return train_data, target_data


def split_k_sets(df):
    k_fold = KFold(shuffle=True, random_state=42)
    return k_fold.split(df)


def parse_data(source, action):
    df = read_csv(filepath_or_buffer=source, na_values="?")
    processed_df = process_na(df=df, action=action)
    encoded_df = encode_columns(df=processed_df)
    test_data, target_data = split_train_target(df=encoded_df)
    return test_data, target_data
