"""Data generator helpers for MEOW-ML."""

import importlib

from meow_ml.src.data_generators.data_generator import DataGenerator
from meow_ml.src.data_generators.two_d_data_generator import TwoDDataGenerator


def import_class(class_path):
    module_path, class_name = class_path.rsplit(".", 1)
    module = importlib.import_module(module_path)
    return getattr(module, class_name)


def _build_data_generator(
    generator_cls,
    data_generator_config,
    split_mode="train",
    scalers=None,
    label_scaler=None,
):
    return generator_cls(
        branch_inputs=data_generator_config["branch_inputs"],
        data_path=data_generator_config["data_path"],
        file_paths=data_generator_config["file_paths"],
        truth=data_generator_config["truth"],
        year1=data_generator_config.get("year1"),
        year2=data_generator_config.get("year2"),
        full_year=data_generator_config.get("full_year", False),
        batch_size=data_generator_config["batch_size"],
        scale_data_method=data_generator_config.get("scale_data_method"),
        randomize_pixels=data_generator_config["randomize_pixels"],
        split_tuple=tuple(data_generator_config["split_tuple"]),
        split_mode=split_mode,
        split_method=data_generator_config["split_method"],
        upscale_params=data_generator_config.get("upscale_params", False),
        test_type=data_generator_config.get("test_type", "cat"),
        bins=data_generator_config.get("bins"),
        scalers=scalers,
        label_scaler=label_scaler,
        scale_y_labels=data_generator_config.get("scale_y_labels", False),
    )


def load_data_generator(data_generator_config, split_mode="train", scalers=None, label_scaler=None):
    generator_cls = DataGenerator
    if data_generator_config.get("class"):
        generator_cls = import_class(data_generator_config["class"])

    return _build_data_generator(
        generator_cls,
        data_generator_config,
        split_mode=split_mode,
        scalers=scalers,
        label_scaler=label_scaler,
    )


def load_2d_data_generator(data_generator_config, split_mode="train", scalers=None):
    return _build_data_generator(
        TwoDDataGenerator,
        data_generator_config,
        split_mode=split_mode,
        scalers=scalers,
        label_scaler=None,
    )
