import os
from findfile import find_files, find_dir

filter_key_words = [
    ".py",
    ".md",
    "readme",
    "log",
    "result",
    "zip",
    ".state_dict",
    ".model",
    ".png",
    "acc_",
    "f1_",
    ".backup",
    ".bak",
]


def detect_infer_dataset(dataset_path, task="apc"):
    dataset_file = []
    if isinstance(dataset_path, str) and os.path.isfile(dataset_path):
        dataset_file.append(dataset_path)
        return dataset_file

    for d in dataset_path:
        if not os.path.exists(d):
            search_path = find_dir(
                os.getcwd(),
                [d, task, "dataset"],
                exclude_key=filter_key_words,
                disable_alert=False,
            )
            dataset_file += find_files(
                search_path,
                [".inference", d],
                exclude_key=["train."] + filter_key_words,
            )
        else:
            dataset_file += find_files(
                d, [".inference", task], exclude_key=["train."] + filter_key_words
            )

    return dataset_file