import os
from pathlib import Path

from datasets import Dataset, DatasetDict, Image

from src.dataset.preprocess import make_nir_transform, make_transform

transform = make_transform(256)
nir_transform = make_nir_transform(256)


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

    data_dir = Path("/data2/C-CUHK/CUHK-CR2")
    nir_dir = Path("/data2/C-CUHK/nir/CUHK-CR2")

    train_cloud = sorted((data_dir / "train/cloud").glob("*.png"))
    train_cloud_nir = sorted((nir_dir / "train/cloud").glob("*.png"))
    train_no_cloud = sorted((data_dir / "train/label").glob("*.png"))
    train_no_cloud_nir = sorted((nir_dir / "train/label").glob("*.png"))
    test_cloud = sorted((data_dir / "test/cloud").glob("*.png"))
    test_cloud_nir = sorted((nir_dir / "test/cloud").glob("*.png"))
    test_no_cloud = sorted((data_dir / "test/label").glob("*.png"))
    test_no_cloud_nir = sorted((nir_dir / "test/label").glob("*.png"))

    dataset = DatasetDict(
        {
            "train": Dataset.from_dict(
                {
                    "cloud": [str(p) for p in train_cloud],
                    "cloud_nir": [str(p) for p in train_cloud_nir],
                    "label": [str(p) for p in train_no_cloud],
                    "label_nir": [str(p) for p in train_no_cloud_nir],
                }
            )
            .cast_column("cloud", Image())
            .cast_column("label", Image())
            .cast_column("cloud_nir", Image())
            .cast_column("label_nir", Image()),
            "test": Dataset.from_dict(
                {
                    "cloud": [str(p) for p in test_cloud],
                    "cloud_nir": [str(p) for p in test_cloud_nir],
                    "label": [str(p) for p in test_no_cloud],
                    "label_nir": [str(p) for p in test_no_cloud_nir],
                }
            )
            .cast_column("cloud", Image())
            .cast_column("label", Image())
            .cast_column("cloud_nir", Image())
            .cast_column("label_nir", 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", "cloud_nir", "gt_nir"])
    dataset.save_to_disk("datasets/CUHK-CR2")

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


def preprocess_function(examples):
    x0_list = []
    x1_list = []
    x0_nir_list = []
    x1_nir_list = []
    for x0_img, x1_img, x0_nir, x1_nir in zip(
        examples["cloud"],
        examples["label"],
        examples["cloud_nir"],
        examples["label_nir"],
    ):
        x0_transformed = transform(x0_img)
        x1_transformed = transform(x1_img)
        x0_nir = nir_transform(x0_nir)
        x1_nir = nir_transform(x1_nir)
        x0_list.append(x0_transformed)
        x1_list.append(x1_transformed)
        x0_nir_list.append(x0_nir)
        x1_nir_list.append(x1_nir)
    return {
        "cloud": x0_list,
        "gt": x1_list,
        "cloud_nir": x0_nir_list,
        "gt_nir": x1_nir_list,
    }
