"""Metadata utilities for the binary coronary segmentation bundle."""

from typing import Dict, Hashable, Mapping, Union

from monai.config.type_definitions import KeysCollection
from monai.data.meta_tensor import MetaTensor
from monai.transforms import MapTransform


class CopyMetadatad(MapTransform):
    """
    Copy MetaData from one MetaTensor to another.

    Args:
        keys: The keys that correspond to the MetaTensors to update.
        reference_key: The key that corresponds to the MetaTensor providing the metadict.
        allow_missing_keys: Whether to allow missing keys.
    """

    def __init__(
        self,
        keys: KeysCollection,
        reference_key: KeysCollection,
        allow_missing_keys: bool = False,
    ) -> None:
        super().__init__(keys, allow_missing_keys)
        self.reference_key = reference_key

    def __call__(self, data: Mapping) -> Dict[Hashable, Union[Dict, MetaTensor]]:
        d = dict(data)
        ref = d[self.reference_key]
        for key in self.key_iterator(d):
            if isinstance(d[key], MetaTensor):
                d[key].meta = ref.meta.copy()
                if hasattr(ref, "applied_operations"):
                    d[key].applied_operations = ref.applied_operations.copy()
        return d