from pandas import DataFrame, read_csv


def replace_values(df) -> DataFrame:
    for column in df.columns:
        df[column].fillna(value=df[column].mean(), inplace=True)
    return df


def process_na(df, action) -> DataFrame:
    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 filter_dataframe(df) -> DataFrame:
    relevant_columns = [
        "HORA",
        "DIASEMANA",
        "COMUNIDAD_AUTONOMA",
        "ISLA",
        "TOT_HERIDOS_LEVES",
        "TOT_HERIDOS_GRAVES",
        "TOT_VEHICULOS_IMPLICADOS",
        "TOT_MUERTOS",
        "TIPO_VIA",
        "LUMINOSIDAD",
        "FACTORES_ATMOSFERICOS",
    ]
    filtered_df = df.filter(items=relevant_columns)
    return filtered_df


def parse_data(source, action) -> DataFrame:
    df = read_csv(filepath_or_buffer=source, na_values="?")
    processed_df = process_na(df=df, action=action)
    return processed_df