from importlib.resources import files
import pandas as pd
import numpy as np

def load_data(filename: str) -> pd.DataFrame:
    filepath = files("bellatrex.datasets") / filename
    return pd.read_csv(filepath)

# Functions to load specific datasets
# Separation between X and y is done 'manually', that is: case by case
def load_binary_data(return_X_y=False):
    if return_X_y is False:
        return load_data('bin_tutorial.csv')
    else:
        df = load_data('bin_tutorial.csv')
        X= df.iloc[:,:-1]
        y = df.iloc[:,-1]
        return X, y

def load_regression_data(return_X_y=False):
    if return_X_y is False:
        return load_data('regress_tutorial.csv')
    else:
        df = load_data('regress_tutorial.csv')
        X= df.iloc[:,:-1]
        y = df.iloc[:,-1]
        return X, y



def load_survival_data(return_X_y=False):
    if return_X_y is False:
        return load_data('surv_tutorial.csv')
    else:
        df = load_data('surv_tutorial.csv')
        X = df.iloc[:,:-2]
        y = df.iloc[:,-2:]

        dtypes_map = {
            str(y.columns[0]): np.bool_,
            str(y.columns[1]): np.float32
            }
        y = y.to_records(index=False, column_dtypes=dtypes_map)

        return X, y

def load_mlc_data(return_X_y=False):
    if return_X_y is False:
        return load_data('mlc_tutorial.csv')
    else:
        df = load_data('mlc_tutorial.csv')
        columns_out = [col for col in df.columns if 'tag' in col]
        columns_in = [col for col in df.columns if 'tag' not in col]

        return df.loc[:,columns_in], df.loc[:,columns_out]


def load_mtr_data(return_X_y=False):
    if return_X_y is False:
        return load_data('mtr_tutorial.csv')
    else:
        df = load_data('mtr_tutorial.csv')
        columns_out = [col for col in df.columns if 'target' in col]
        columns_in = [col for col in df.columns if 'target' not in col]

        return df.loc[:,columns_in], df.loc[:,columns_out]
