model.py

import os
import joblib
from pathlib import Path

import imblearn


def load_models(model_name, target_column):
    """Load models from disk"""
    model_root = Path("./models")
    target_column_root = model_root / target_column

    if target_column == "OfficerRace":
        model_selected_root = target_column_root / model_name
        models = [joblib.load(f) for f in model_selected_root.iterdir()]
        names = [path.stem.split("_")[-1] for path in model_selected_root.iterdir()]
        return names, models

    if target_column == "OfficerGender":
        model_path = target_column_root / f"{model_name}.pkl"
        return joblib.load(model_path)

    if target_column == "PenaltyCat":
        model_path = target_column_root / f"{model_name}.pkl"
        return joblib.load(model_path)


def get_folder_names(directory):
    return [name for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))]


def available_models(target_column):
    """Return available models for a given target column"""

    if target_column == "OfficerRace":
        return get_folder_names("models/OfficerRace")

    elif target_column == "OfficerGender":
        return [os.path.splitext(name)[0] for name in os.listdir(f"models/{target_column}") if name.endswith(".pkl")]

    elif target_column == "PenaltyCat":
        return [os.path.splitext(name)[0] for name in os.listdir(f"models/{target_column}") if name != ("scaler_model"
                                                                                                        ".pkl")]