from core.datasets.BaseDataset import BaseDataset
from core.datasets.classic.BreastCancerDataset import BreastCancerDataset
from core.datasets.classic.DiabetesDataset import DiabetesDataset
from core.datasets.classic.GlassDataset import GlassDataset
from core.datasets.classic.IrisDataset import IrisDataset
from core.datasets.classic.LeafDataset import LeafDataset
from core.datasets.classic.ParkinsonsDataset import ParkinsonsDataset
from core.datasets.classic.WineDataset import WineDataset
from core.datasets.custom.CustomNDataset import CustomNDataset
from core.datasets.image.DigitsDataset import DigitsDataset
from core.datasets.image.MnistDataset import MnistDataset
from functools import lru_cache


@lru_cache(maxsize=50)
def create_dataset(name, custom_dim, seed) -> BaseDataset:
    """Create a dataset with LRU caching based on name, custom_dim, and seed."""
    if name == "breast_cancer":
        return BreastCancerDataset(seed)
    if name == "custom":
        if custom_dim is None:
            raise ValueError("'custom_dim' must be provided for custom dataset.")
        return CustomNDataset(custom_dim, seed)
    if name == "diabetes":
        return DiabetesDataset(seed)
    if name == "digits":
        return DigitsDataset(seed)
    if name == "glass":
        return GlassDataset(seed)
    if name == "iris":
        return IrisDataset(seed)
    if name == "leaf":
        return LeafDataset(seed)
    if name == "mnist":
        return MnistDataset(seed)
    if name == "parkinsons":
        return ParkinsonsDataset(seed)
    if name == "wine":
        return WineDataset(seed)

    raise ValueError(f"Unknown dataset: {name}")


def clear_dataset_cache():
    """Clear the dataset cache. Useful for memory management."""
    create_dataset.cache_clear()
