# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import os
import tempfile
import unittest
from pathlib import Path

import numpy as np
import torch
from parameterized import parameterized

from monai.config import PathLike
from monai.data.folder_layout import FolderLayoutBase
from monai.data.meta_tensor import MetaTensor
from monai.transforms import SaveImaged
from monai.utils import optional_import

_, has_itk = optional_import("itk", allow_namespace_pkg=True)
nib, has_nibabel = optional_import("nibabel")

TEST_CASE_1 = [
    {"img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nii.gz"})},
    ".nii.gz",
    False,
]

TEST_CASE_2 = [
    {
        "img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nii.gz"}),
        "patch_index": 6,
    },
    ".nii.gz",
    False,
]

TEST_CASE_3 = [
    {
        "img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nrrd"}),
        "patch_index": 6,
    },
    ".nrrd",
    False,
]


@unittest.skipUnless(has_itk, "itk not installed")
class TestSaveImaged(unittest.TestCase):

    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
    def test_saved_content(self, test_data, output_ext, resample):
        with tempfile.TemporaryDirectory() as tempdir:
            trans = SaveImaged(
                keys=["img", "pred"],
                output_dir=tempdir,
                output_ext=output_ext,
                resample=resample,
                allow_missing_keys=True,
            )
            trans(test_data)

            patch_index = test_data["img"].meta.get("patch_index", None)
            patch_index = f"_{patch_index}" if patch_index is not None else ""
            filepath = os.path.join("testfile0", "testfile0" + "_trans" + patch_index + output_ext)
            self.assertTrue(os.path.exists(os.path.join(tempdir, filepath)))

    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
    def test_custom_folderlayout(self, test_data, output_ext, resample):

        class TestFolderLayout(FolderLayoutBase):

            def __init__(self, basepath: Path, extension: str, makedirs: bool):
                self.basepath = basepath
                self.ext = extension
                self.makedirs = makedirs

            def filename(self, **kwargs) -> PathLike:
                p = self.basepath / str(kwargs["subdirectory"])
                if not p.exists() and self.makedirs:
                    p.mkdir()

                return p / (str(kwargs["filename"]) + self.ext)

        def name_formatter(metadict: dict, _) -> dict:
            # "[filename].[ext]"
            # quick and dirty split on .
            base_filename = metadict["filename_or_obj"].split(".")[0]

            return {"subdirectory": base_filename, "filename": "image"}

        with tempfile.TemporaryDirectory() as tempdir:
            trans = SaveImaged(
                keys=["img", "pred"],
                resample=resample,
                allow_missing_keys=True,
                output_name_formatter=name_formatter,
                folder_layout=TestFolderLayout(basepath=Path(tempdir), extension=output_ext, makedirs=True),
            )
            trans(test_data)

            filepath = os.path.join("testfile0", "image" + output_ext)
            self.assertTrue(os.path.exists(os.path.join(tempdir, filepath)))

    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
    def test_includes_metadata(self, test_data, output_ext, resample):
        with tempfile.TemporaryDirectory() as tempdir:
            trans = SaveImaged(
                keys=["img", "pred"],
                output_dir=tempdir,
                output_ext=output_ext,
                resample=resample,
                allow_missing_keys=True,
                savepath_in_metadict=True,
            )
            trans(test_data)

            self.assertTrue("saved_to" in test_data["img"].meta.keys())
            self.assertTrue(os.path.exists(test_data["img"].meta["saved_to"]))


@unittest.skipUnless(has_nibabel, "nibabel not installed")
class TestSaveImagedChannelDim(unittest.TestCase):

    @parameterized.expand(
        [
            ("image_first", ("image", "label"), False),
            ("label_first", ("label", "image"), False),
            ("metatensor", ("image", "label"), True),
        ]
    )
    def test_mixed_channel_dimensions(self, _, keys, use_metatensor):
        image = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
        label = (image + 100)[None]
        with tempfile.TemporaryDirectory() as tempdir:
            saver = SaveImaged(
                keys=keys,
                output_dir=tempdir,
                writer="NibabelWriter",
                resample=False,
                separate_folder=False,
                print_log=False,
            )
            data = {}
            for key, array in (("image", image), ("label", label)):
                metadata = {"filename_or_obj": key, "spatial_shape": (2, 3, 4)}
                data[key] = MetaTensor(array, meta=metadata) if use_metatensor else array
                if not use_metatensor:
                    data[f"{key}_meta_dict"] = metadata
            for _ in range(2):
                saver(data)
                for key, expected in (("image", image), ("label", label[0])):
                    saved = nib.load(os.path.join(tempdir, f"{key}_trans.nii.gz")).get_fdata()
                    np.testing.assert_array_equal(saved, expected)


if __name__ == "__main__":
    unittest.main()