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

def load_data_generator(data_generator_config, split_mode='train', scalers=None, label_scaler=None):

    data_generator = DataGenerator(
        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", None),
        year2=data_generator_config.get("year2", None),
        # full_hyperspectral=data_generator_config.get("full_hyperspectral", None),
        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", None),
        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", None),
        scalers=scalers,
        label_scaler=label_scaler, # TODO: delete if y scaling not useful
        scale_y_labels=data_generator_config.get("scale_y_labels", False),
    )

    return data_generator


def load_2d_data_generator(data_generator_config, split_mode='train', scalers=None):

    data_generator = TwoDDataGenerator(
        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", None),
        year2=data_generator_config.get("year2", None),
        # full_hyperspectral=data_generator_config.get("full_hyperspectral", None),
        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", None),
        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", None),
        scalers=scalers
    )

    return data_generator
