from os import path as osp

import msgspec
from traiNNer.data.paired_image_dataset import PairedImageDataset
from traiNNer.utils.redux_options import DatasetOptions


def test_pairedimagedataset() -> None:
    """Test dataset: PairedImageDataset"""

    opt_str = r"""
name: Test
type: PairedImageDataset
dataroot_gt: [
datasets/val/dataset1/hr,
datasets/val/dataset1/hr2,
]
dataroot_lq: [
datasets/val/dataset1/lr,
datasets/val/dataset1/lr2
]
filename_tmpl: '{}'
io_backend:
    type: disk

scale: 4
mean: [0.5, 0.5, 0.5]
std: [0.5, 0.5, 0.5]
gt_size: 128
use_hflip: true
use_rot: true

phase: train
"""
    image_names = (7, 8, 9)
    opt = msgspec.yaml.decode(opt_str, type=DatasetOptions, strict=True)

    dataset = PairedImageDataset(opt)
    assert dataset.io_backend_opt["type"] == "disk"  # io backend
    assert len(dataset) == 3  # whether to read correct meta info
    assert dataset.mean == [0.5, 0.5, 0.5]

    # ------------------ test scan folder mode -------------------- #
    opt.io_backend = {"type": "disk"}
    dataset = PairedImageDataset(opt)
    assert dataset.io_backend_opt["type"] == "disk"  # io backend
    assert len(dataset) == 3  # whether to correctly scan folders

    # test __getitem__
    result = dataset.__getitem__(0)
    # check returned keys
    expected_keys = ["lq", "gt", "lq_path", "gt_path"]
    assert set(expected_keys).issubset(set(result.keys()))
    # check shape and contents
    assert (
        "gt" in result
        and "lq" in result
        and "lq_path" in result
        and "gt_path" in result
    )
    assert result["gt"].shape == (3, 128, 128)
    assert result["lq"].shape == (3, 32, 32)
    assert osp.normpath(result["lq_path"]) in {
        osp.normpath(f"datasets/val/dataset1/lr/{x:04d}.png") for x in image_names
    }
    assert osp.normpath(result["gt_path"]) in {
        osp.normpath(f"datasets/val/dataset1/hr/{x:04d}.png") for x in image_names
    }

    # ------------------ test lmdb backend and with y channel-------------------- #
    # TODO
    # opt["dataroot_gt"] = "tests/data/gt.lmdb"
    # opt["dataroot_lq"] = "tests/data/lq.lmdb"
    # opt["io_backend"] = {"type": "lmdb"}
    # opt["color"] = "y"
    # opt["mean"] = [0.5]
    # opt["std"] = [0.5]

    # dataset = PairedImageDataset(opt)
    # assert dataset.io_backend_opt["type"] == "lmdb"  # io backend
    # assert len(dataset) == 2  # whether to read correct meta info
    # assert dataset.std == [0.5]

    # # test __getitem__
    # result = dataset.__getitem__(1)
    # # check returned keys
    # expected_keys = ["lq", "gt", "lq_path", "gt_path"]
    # assert set(expected_keys).issubset(set(result.keys()))
    # # check shape and contents
    # assert (
    #     "gt" in result
    #     and "lq" in result
    #     and "lq_path" in result
    #     and "gt_path" in result
    # )
    # assert result["gt"].shape == (1, 128, 128)
    # assert result["lq"].shape == (1, 32, 32)
    # assert result["lq_path"] == "comic"
    # assert result["gt_path"] == "comic"

    # ------------------ test case: val/test mode -------------------- #
    # TODO
    # opt["phase"] = "test"
    # opt["io_backend"] = {"type": "lmdb"}
    # dataset = PairedImageDataset(opt)

    # # test __getitem__
    # result = dataset.__getitem__(0)
    # # check returned keys
    # expected_keys = ["lq", "gt", "lq_path", "gt_path"]
    # assert set(expected_keys).issubset(set(result.keys()))
    # # check shape and contents
    # assert (
    #     "gt" in result
    #     and "lq" in result
    #     and "lq_path" in result
    #     and "gt_path" in result
    # )
    # assert result["gt"].shape == (1, 480, 492)
    # assert result["lq"].shape == (1, 120, 123)
    # assert result["lq_path"] == "baboon"
    # assert result["gt_path"] == "baboon"