import nibabel as nib
import torch

from monai.data import DataLoader, Dataset, ImageDataset, ZipDataset
from monai.data.utils import worker_init_fn
from monai.transforms import AddChannel, Compose, EnsureType, RandGaussianSmooth, RandSpatialCrop, ScaleIntensity

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def main():
    print("device", device)
    # create 4 (3x3) test images
    test_img1 = torch.Tensor(
        [
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0],
        ]
    )
    test_img2 = test_img1 * 2.0
    test_img3 = test_img1 * 3.0
    test_img4 = test_img1 * 4.0
    test_img1 = nib.Nifti1Image(test_img1.numpy(), None)
    test_img2 = nib.Nifti1Image(test_img2.numpy(), None)
    test_img3 = nib.Nifti1Image(test_img3.numpy(), None)
    test_img4 = nib.Nifti1Image(test_img4.numpy(), None)
    # write images
    nib.save(test_img1, "test_img1.nii.gz")
    nib.save(test_img2, "tes_img3.nii.gz")
    nib.save(test_img2, "test_img4.nii.gz")
    ds_transforms = Compose(
        [
            AddChannel(),
            RandGaussianSmooth(sigma_x=(0, 1), sigma_y=(0, 1), sigma_z=(0, 1), prob=1),
            # RandRotate((-180,180),prob=1.),
            RandSpatialCrop((2, 2), random_center=True, random_size=False),
            EnsureType(),
        ]
    )
    ds1 = ImageDataset(
        image_files=[
            "test_img1.nii.gz",
            "test_img2.nii.gz",
            "test_img3.nii.gz",
            "test_img4.nii.gz",
        ],
        labels=[1, 2, 3, 4],
        transform=ds_transforms,
        reader="NibabelReader",
    )
    ds2 = Dataset(data=[1, 2, 3, 4])
    ds = ZipDataset([ds1, ds2])
    ds_loader = DataLoader(
        ds,
        batch_size=2,
        num_workers=4,
        worker_init_fn=worker_init_fn,
    )
    for epoch in range(0, 5):
        print("Epoch [{0}]: \t".format(epoch))
        train(ds_loader)


def train(train_loader):
    for i, (input, target, another_label) in enumerate(train_loader):
        if i == 0:
            image_0 = input[0, 0, :, :].detach().cpu().numpy()
            print("image0:", image_0)
    return None


if __name__ == "__main__":
    main()