import os
from pathlib import Path

from datasets import Dataset, DatasetDict, Image
from src.dataset.preprocess import make_transform

transform = make_transform(512)


def get_dataset() -> tuple[Dataset, Dataset]:
    if os.path.exists("datasets/CUHK-CR1"):
        dataset = DatasetDict.load_from_disk("datasets/CUHK-CR1")
        return dataset["train"], dataset["test"]

    data_dir = Path("/data2/C-CUHK/CUHK-CR1")

    train_cloud = sorted((data_dir / "train/cloud").glob("*.png"))
    train_no_cloud = sorted((data_dir / "train/label").glob("*.png"))
    test_cloud = sorted((data_dir / "test/cloud").glob("*.png"))
    test_no_cloud = sorted((data_dir / "test/label").glob("*.png"))

    dataset = DatasetDict(
        {
            "train": Dataset.from_dict(
                {
                    "cloud": [str(p) for p in train_cloud],
                    "label": [str(p) for p in train_no_cloud],
                }
            )
            .cast_column("cloud", Image())
            .cast_column("label", Image()),
            "test": Dataset.from_dict(
                {
                    "cloud": [str(p) for p in test_cloud],
                    "label": [str(p) for p in test_no_cloud],
                }
            )
            .cast_column("cloud", Image())
            .cast_column("label", Image()),
        }
    )
    dataset = dataset.map(
        preprocess_function,
        batched=True,
        batch_size=32,
        remove_columns=dataset["train"].column_names,
    )
    dataset.set_format(type="torch", columns=["cloud", "gt"])
    dataset.save_to_disk("datasets/CUHK-CR1")

    return dataset["train"], dataset["test"]


def preprocess_function(examples):
    x0_list = []
    x1_list = []
    for x0_img, x1_img in zip(examples["cloud"], examples["label"]):
        x0_transformed = transform(x0_img)
        x1_transformed = transform(x1_img)
        x0_list.append(x0_transformed)
        x1_list.append(x1_transformed)
    return {"cloud": x0_list, "gt": x1_list}
