# 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 unittest import skipUnless

import numpy as np
from parameterized import parameterized

from monai.bundle import ConfigParser, ReferenceResolver
from monai.bundle.config_item import ConfigItem
from monai.data import DataLoader, Dataset
from monai.transforms import Compose, LoadImaged, RandTorchVisiond
from monai.utils import min_version, optional_import
from tests.utils import TimedCall

_, has_tv = optional_import("torchvision", "0.8.0", min_version)


@TimedCall(seconds=100, force_quit=True)
def case_pdb(sarg=None):
    config = {"transform": {"_target_": "Compose", "transforms": [], "_debug_": True}}
    parser = ConfigParser(config=config)
    parser.get_parsed_content()


@TimedCall(seconds=100, force_quit=True)
def case_pdb_inst(sarg=None):
    config = {"transform": {"_target_": "Compose", "transforms": [], "_mode_": "debug"}}
    parser = ConfigParser(config=config)
    return parser.transform


# test the resolved and parsed instances
TEST_CASE_1 = [
    {
        "transform": {
            "_target_": "Compose",
            "transforms": [
                {"_target_": "LoadImaged", "keys": "image"},
                # test relative id in `keys`
                {"_target_": "RandTorchVisiond", "keys": "@##0#keys", "name": "ColorJitter", "brightness": 0.25},
            ],
        },
        "dataset": {"_target_": "Dataset", "data": [1, 2], "transform": "@transform"},
        "dataloader": {
            "_target_": "DataLoader",
            # test relative id in `dataset`
            "dataset": "@##dataset",
            "batch_size": 2,
            "collate_fn": "$monai.data.list_data_collate",
        },
    },
    ["transform", "transform#transforms#0", "transform#transforms#1", "dataset", "dataloader"],
    [Compose, LoadImaged, RandTorchVisiond, Dataset, DataLoader],
]


class TestClass:
    @staticmethod
    def compute(a, b, func=lambda x, y: x + y):
        return func(a, b)

    @classmethod
    def cls_compute(cls, a, b, func=lambda x, y: x + y):
        return cls.compute(a, b, func)

    def __call__(self, a, b):
        return self.compute(a, b)


TEST_CASE_2 = [
    {
        "basic_func": "$lambda x, y: x + y",
        "static_func": "$TestClass.compute",
        "cls_func": "$TestClass.cls_compute",
        "lambda_static_func": "$lambda x, y: TestClass.compute(x, y)",
        "lambda_cls_func": "$lambda x, y: TestClass.cls_compute(x, y)",
        "compute": {"_target_": "tests.test_config_parser.TestClass.compute", "func": "@basic_func"},
        "cls_compute": {"_target_": "tests.test_config_parser.TestClass.cls_compute", "func": "@basic_func"},
        "call_compute": {"_target_": "tests.test_config_parser.TestClass"},
        "error_func": "$TestClass.__call__",
        "<test>": "$lambda x, y: x + y",
    }
]

TEST_CASE_3 = [
    {
        "A": 1,
        "B": "@A",
        "C": "@#A",
        "D": {"key": "@##A", "value1": 2, "value2": "%#value1", "value3": [3, 4, "@#1", "$100 + @#0 + @##value1"]},
    }
]

TEST_CASE_4 = [{"A": 1, "B": "@A", "C": "@D", "E": "$'test' + '@F'"}]

TEST_CASE_5 = [{"training": {"A": 1, "A_B": 2}, "total": "$@training#A + @training#A_B + 1"}, 4]


class TestConfigParser(unittest.TestCase):
    def test_config_content(self):
        test_config = {"preprocessing": [{"_target_": "LoadImage"}], "dataset": {"_target_": "Dataset"}}
        parser = ConfigParser(config=test_config)
        # test `get`, `set`, `__getitem__`, `__setitem__`
        self.assertEqual(str(parser.get()), str(test_config))
        parser.set(config=test_config)
        self.assertListEqual(parser["preprocessing"], test_config["preprocessing"])
        parser["dataset"] = {"_target_": "CacheDataset"}
        self.assertEqual(parser["dataset"]["_target_"], "CacheDataset")
        # test nested ids
        parser["dataset#_target_"] = "Dataset"
        self.assertEqual(parser["dataset#_target_"], "Dataset")
        parser.update({"dataset#_target_1": "Dataset1"})
        self.assertEqual(parser["dataset#_target_1"], "Dataset1")
        # test int id
        parser.set(["test1", "test2", "test3"])
        parser[1] = "test4"
        self.assertEqual(parser[1], "test4")

    @parameterized.expand([TEST_CASE_1])
    @skipUnless(has_tv, "Requires torchvision >= 0.8.0.")
    def test_parse(self, config, expected_ids, output_types):
        parser = ConfigParser(config=config, globals={"monai": "monai"})
        # test lazy instantiation with original config content
        parser["transform"]["transforms"][0]["keys"] = "label1"
        trans = parser.get_parsed_content(id="transform#transforms#0")
        self.assertEqual(trans.keys[0], "label1")
        # test re-use the parsed content or not with the `lazy` option
        self.assertEqual(trans, parser.get_parsed_content(id="transform#transforms#0"))
        self.assertEqual(trans, parser.get_parsed_content(id="transform#transforms#0", lazy=True))
        self.assertNotEqual(trans, parser.get_parsed_content(id="transform#transforms#0", lazy=False))
        # test new nested id
        parser.set("fake_key", "transform#other_transforms#keys", True)
        self.assertEqual(parser.get(id="transform#other_transforms#keys"), "fake_key")
        # remove temp fake data
        parser["transform"].pop("other_transforms")
        # test update nested id
        parser["transform#transforms#0#keys"] = "label2"
        self.assertEqual(parser.get_parsed_content(id="transform#transforms#0").keys[0], "label2")

        for id, cls in zip(expected_ids, output_types):
            self.assertTrue(isinstance(parser.get_parsed_content(id), cls))
        # test root content
        root = parser.get_parsed_content(id="")
        for v, cls in zip(root.values(), [Compose, Dataset, DataLoader]):
            self.assertTrue(isinstance(v, cls))
        # test default value
        self.assertEqual(parser.get_parsed_content(id="abc", default=ConfigItem(12345, "abc")), 12345)

    @parameterized.expand([TEST_CASE_2])
    def test_function(self, config):
        parser = ConfigParser(config=config, globals={"TestClass": TestClass})
        for id in config:
            if id in ("compute", "cls_compute"):
                parser[f"{id}#_mode_"] = "partial"
            func = parser.get_parsed_content(id=id)
            self.assertTrue(id in parser.ref_resolver.resolved_content)
            if id == "error_func":
                with self.assertRaises(TypeError):
                    func(1, 2)
                continue
            self.assertEqual(func(1, 2), 3)

    @parameterized.expand([TEST_CASE_3])
    def test_relative_id(self, config):
        parser = ConfigParser(config=config)
        for id in config:
            item = parser.get_parsed_content(id=id)
            if isinstance(item, int):
                self.assertEqual(item, 1)
            if isinstance(item, dict):
                self.assertEqual(str(item), str({"key": 1, "value1": 2, "value2": 2, "value3": [3, 4, 4, 105]}))

    def test_macro_replace(self):
        with tempfile.TemporaryDirectory() as tempdir:
            another_file = os.path.join(tempdir, "another.json")
            ConfigParser.export_config_file(config={"E": 4}, filepath=another_file)
            # test macro with id, relative id, and macro in another file
            config = {"A": {"B": 1, "C": 2}, "D": [3, "%A#B", "%#0", f"%{another_file}#E"]}
            parser = ConfigParser(config=config)
            parser.resolve_macro_and_relative_ids()
            self.assertEqual(str(parser.get()), str({"A": {"B": 1, "C": 2}, "D": [3, 1, 3, 4]}))

    @parameterized.expand([TEST_CASE_4])
    def test_allow_missing_reference(self, config):
        default = ReferenceResolver.allow_missing_reference
        ReferenceResolver.allow_missing_reference = True
        parser = ConfigParser(config=config)

        for id in config:
            item = parser.get_parsed_content(id=id)
            if id in ("A", "B"):
                self.assertEqual(item, 1)
            elif id == "C":
                self.assertEqual(item, "@D")
            elif id == "E":
                self.assertEqual(item, "test@F")

        # restore the default value
        ReferenceResolver.allow_missing_reference = default
        with self.assertRaises(ValueError):
            parser.parse()
            parser.get_parsed_content(id="E")

    def test_list_expressions(self):
        config = {
            "transform": {
                "_target_": "Compose",
                "transforms": [{"_target_": "RandScaleIntensity", "factors": 0.5, "prob": 1.0}],
            },
            "training": ["$monai.utils.set_determinism(seed=123)", "$@transform(np.asarray([1, 2]))"],
        }
        parser = ConfigParser(config=config)
        parser.get_parsed_content("training", lazy=True, instantiate=True, eval_expr=True)
        np.testing.assert_allclose(parser.get_parsed_content("training#1", lazy=True), [0.7942, 1.5885], atol=1e-4)

    def test_contains(self):
        empty_parser = ConfigParser({})
        empty_parser.parse()

        parser = ConfigParser({"value": 1, "entry": "string content", "array": [1, 2]})
        parser.parse()

        with self.subTest("Testing empty parser"):
            self.assertFalse("something" in empty_parser)
        with self.assertRaises(KeyError):
            empty_parser["something"]
        empty_parser["osmething"] = "test"
        with self.assertRaises(KeyError):
            empty_parser["something"]

        with self.subTest("Testing with keys"):
            self.assertTrue("value" in parser)
            self.assertFalse("value1" in parser)
            self.assertTrue("entry" in parser)
            self.assertFalse("entr" in parser)
            self.assertFalse("array#2" in parser)

    def test_lambda_reference(self):
        configs = {
            "patch_size": [8, 8],
            "transform": {"_target_": "Lambda", "func": "$lambda x: x.reshape((1, *@patch_size))"},
        }
        parser = ConfigParser(config=configs)
        trans = parser.get_parsed_content(id="transform")
        result = trans(np.ones(64))
        self.assertTupleEqual(result.shape, (1, 8, 8))

    def test_non_str_target(self):
        configs = {
            "fwd": {"_target_": "$@model.forward", "x": "$torch.rand(1, 3, 256, 256)", "_mode_": "partial"},
            "model": {"_target_": "monai.networks.nets.resnet.resnet18", "pretrained": False, "spatial_dims": 2},
        }
        self.assertTrue(callable(ConfigParser(config=configs).fwd))
        self.assertTupleEqual(tuple(ConfigParser(config=configs).fwd().shape), (1, 400))

    def test_error_instance(self):
        config = {"transform": {"_target_": "Compose", "transforms_wrong_key": []}}
        parser = ConfigParser(config=config)
        with self.assertRaises(RuntimeError):
            parser.get_parsed_content("transform", instantiate=True, eval_expr=True)

    def test_pdb(self):
        with self.assertRaisesRegex(RuntimeError, ".*bdb.BdbQuit.*"):
            case_pdb()
        self.assertEqual(case_pdb_inst(), None)  # pdb.runcall without input is None

    def test_get_via_attributes(self):
        config = {
            "A": {"B": {"C": 1}},
            "my_dims": 2,
            "dims_1": "$@my_dims + 1",
            "patch_size": [8, 8],
            "transform": {"_target_": "Lambda", "func": "$lambda x: x.reshape((1, *@patch_size))"},
        }
        parser = ConfigParser(config=config)
        self.assertEqual(parser.A, {"B": {"C": 1}})
        self.assertEqual(parser.dims_1, 3)

        trans = parser.transform
        result = trans(np.ones(64))
        self.assertTupleEqual(result.shape, (1, 8, 8))

    def test_builtin(self):
        config = {"import statements": "$import math", "calc": {"_target_": "math.isclose", "a": 0.001, "b": 0.001}}
        self.assertEqual(ConfigParser(config).calc, True)

    @parameterized.expand([TEST_CASE_5])
    def test_substring_reference(self, config, expected):
        parser = ConfigParser(config=config)
        self.assertEqual(parser.get_parsed_content("total"), expected)


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