# Copyright The PyTorch Lightning team.
#
# 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 functools import partial
from types import FunctionType
from typing import Any, Callable, Dict, List, Optional, Union

from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.utilities.exceptions import MisconfigurationException

_REGISTERED_FUNCTION = Dict[str, Any]


class FlashRegistry:
    """This class is used to register function or :class:`functools.partial` class to a registry."""

    def __init__(self, name: str, verbose: bool = False) -> None:
        self.name = name
        self.functions: List[_REGISTERED_FUNCTION] = []
        self._verbose = verbose

    def __len__(self) -> int:
        return len(self.functions)

    def __contains__(self, key) -> bool:
        return any(key == e["name"] for e in self.functions)

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(name={self.name}, functions={self.functions})'

    def get(
        self,
        key: str,
        with_metadata: bool = False,
        strict: bool = True,
        **metadata,
    ) -> Union[Callable, _REGISTERED_FUNCTION, List[_REGISTERED_FUNCTION], List[Callable]]:
        """
        This function is used to gather matches from the registry:

        Args:
            key: Name of the registered function.
            with_metadata: Whether to include the associated metadata in the return value.
            strict: Whether to return all matches or just one.
            metadata: Metadata used to filter against existing registry item's metadata.
        """
        matches = [e for e in self.functions if key == e["name"]]
        if not matches:
            raise KeyError(f"Key: {key} is not in {repr(self)}")

        if metadata:
            matches = [m for m in matches if metadata.items() <= m["metadata"].items()]
            if not matches:
                raise KeyError("Found no matches that fit your metadata criteria. Try removing some metadata")

        matches = [e if with_metadata else e["fn"] for e in matches]
        return matches[0] if strict else matches

    def remove(self, key: str) -> None:
        self.functions = [f for f in self.functions if f["name"] != key]

    def _register_function(
        self,
        fn: Callable,
        name: Optional[str] = None,
        override: bool = False,
        metadata: Optional[Dict[str, Any]] = None
    ):
        if not isinstance(fn, FunctionType) and not isinstance(fn, partial):
            raise MisconfigurationException(f"You can only register a function, found: {fn}")

        name = name or fn.__name__

        if self._verbose:
            rank_zero_info(f"Registering: {fn.__name__} function with name: {name} and metadata: {metadata}")

        item = {"fn": fn, "name": name, "metadata": metadata or {}}

        matching_index = self._find_matching_index(item)
        if override and matching_index is not None:
            self.functions[matching_index] = item
        else:
            if matching_index is not None:
                raise MisconfigurationException(
                    f"Function with name: {name} and metadata: {metadata} is already present within {self}."
                    " HINT: Use `override=True`."
                )
            self.functions.append(item)

    def _find_matching_index(self, item: _REGISTERED_FUNCTION) -> Optional[int]:
        for idx, fn in enumerate(self.functions):
            if all(fn[k] == item[k] for k in ("fn", "name", "metadata")):
                return idx

    def __call__(
        self,
        fn: Optional[Callable[..., Any]] = None,
        name: Optional[str] = None,
        override: bool = False,
        **metadata
    ) -> Callable:
        """
        This function is used to register new functions to the registry along their metadata.

        Functions can be filtered using metadata using the ``get`` function.

        """
        if fn is not None:
            self._register_function(fn=fn, name=name, override=override, metadata=metadata)
            return fn

        # raise the error ahead of time
        if not (name is None or isinstance(name, str)):
            raise TypeError(f'`name` must be a str, found {name}')

        def _register(cls):
            self._register_function(fn=cls, name=name, override=override, metadata=metadata)
            return cls

        return _register

    def available_keys(self) -> List[str]:
        return sorted(v["name"] for v in self.functions)

flash/core/registry.py
# Copyright The PyTorch Lightning team.
#
# 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 functools import partial
from types import FunctionType
from typing import Any, Callable, Dict, List, Optional, Union

from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.utilities.exceptions import MisconfigurationException

_REGISTERED_FUNCTION = Dict[str, Any]


class FlashRegistry:
    """This class is used to register function or :class:`functools.partial` class to a registry."""

    def __init__(self, name: str, verbose: bool = False) -> None:
        self.name = name
        self.functions: List[_REGISTERED_FUNCTION] = []
        self._verbose = verbose

    def __len__(self) -> int:
        return len(self.functions)

    def __contains__(self, key) -> bool:
        return any(key == e["name"] for e in self.functions)

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(name={self.name}, functions={self.functions})'

    def get(
        self,
        key: str,
        with_metadata: bool = False,
        strict: bool = True,
        **metadata,
    ) -> Union[Callable, _REGISTERED_FUNCTION, List[_REGISTERED_FUNCTION], List[Callable]]:
        """
        This function is used to gather matches from the registry:

        Args:
            key: Name of the registered function.
            with_metadata: Whether to include the associated metadata in the return value.
            strict: Whether to return all matches or just one.
            metadata: Metadata used to filter against existing registry item's metadata.
        """
        matches = [e for e in self.functions if key == e["name"]]
        if not matches:
            raise KeyError(f"Key: {key} is not in {type(self).__name__}")

        if metadata:
            matches = [m for m in matches if metadata.items() <= m["metadata"].items()]
            if not matches:
                raise KeyError("Found no matches that fit your metadata criteria. Try removing some metadata")

        matches = [e if with_metadata else e["fn"] for e in matches]
        return matches[0] if strict else matches

    def remove(self, key: str) -> None:
        self.functions = [f for f in self.functions if f["name"] != key]

    def _register_function(
        self,
        fn: Callable,
        name: Optional[str] = None,
        override: bool = False,
        metadata: Optional[Dict[str, Any]] = None
    ):
        if not isinstance(fn, FunctionType) and not isinstance(fn, partial):
            raise MisconfigurationException(f"You can only register a function, found: {fn}")

        name = name or fn.__name__

        if self._verbose:
            rank_zero_info(f"Registering: {fn.__name__} function with name: {name} and metadata: {metadata}")

        item = {"fn": fn, "name": name, "metadata": metadata or {}}

        matching_index = self._find_matching_index(item)
        if override and matching_index is not None:
            self.functions[matching_index] = item
        else:
            if matching_index is not None:
                raise MisconfigurationException(
                    f"Function with name: {name} and metadata: {metadata} is already present within {self}."
                    " HINT: Use `override=True`."
                )
            self.functions.append(item)

    def _find_matching_index(self, item: _REGISTERED_FUNCTION) -> Optional[int]:
        for idx, fn in enumerate(self.functions):
            if all(fn[k] == item[k] for k in ("fn", "name", "metadata")):
                return idx

    def __call__(
        self,
        fn: Optional[Callable[..., Any]] = None,
        name: Optional[str] = None,
        override: bool = False,
        **metadata
    ) -> Callable:
        """
        This function is used to register new functions to the registry along their metadata.

        Functions can be filtered using metadata using the ``get`` function.

        """
        if fn is not None:
            self._register_function(fn=fn, name=name, override=override, metadata=metadata)
            return fn

        # raise the error ahead of time
        if not (name is None or isinstance(name, str)):
            raise TypeError(f'`name` must be a str, found {name}')

        def _register(cls):
            self._register_function(fn=cls, name=name, override=override, metadata=metadata)
            return cls

        return _register

    def available_keys(self) -> List[str]:
        return sorted(v["name"] for v in self.functions)

flash/core/registry.py
diff --git a/flash/core/registry.py b/flash/core/registry.py
index ff3c99c336..61794424ce 100644
--- a/flash/core/registry.py
+++ b/flash/core/registry.py
@@ -56,7 +56,7 @@ def get(
         """
         matches = [e for e in self.functions if key == e["name"]]
         if not matches:
-            raise KeyError(f"Key: {key} is not in {repr(self)}")
+            raise KeyError(f"Key: {key} is not in {type(self).__name__}")
 
         if metadata:
             matches = [m for m in matches if metadata.items() <= m["metadata"].items()]
diff --git a/tests/core/test_registry.py b/tests/core/test_registry.py
index 061c6f4504..3af891aa3a 100644
--- a/tests/core/test_registry.py
+++ b/tests/core/test_registry.py
@@ -28,19 +28,19 @@ def my_model(nc_input=5, nc_output=6):
         return nn.Linear(nc_input, nc_output), nc_input, nc_output
 
     with pytest.raises(MisconfigurationException, match="You can only register a function, found: Linear"):
-        backbones(nn.Linear(1, 1), name="cho")
+        backbones(nn.Linear(1, 1), name="foo")
 
-    backbones(my_model, name="cho", override=True)
+    backbones(my_model, name="foo", override=True)
 
-    with pytest.raises(MisconfigurationException, match="Function with name: cho and metadata: {}"):
-        backbones(my_model, name="cho", override=False)
+    with pytest.raises(MisconfigurationException, match="Function with name: foo and metadata: {}"):
+        backbones(my_model, name="foo", override=False)
 
     with pytest.raises(KeyError, match="Found no matches"):
-        backbones.get("cho", foo="bar")
+        backbones.get("foo", baz="bar")
 
-    backbones.remove("cho")
-    with pytest.raises(KeyError, match="Key: cho is not in FlashRegistry"):
-        backbones.get("cho")
+    backbones.remove("foo")
+    with pytest.raises(KeyError, match="Key: foo is not in FlashRegistry"):
+        backbones.get("foo")
 
     with pytest.raises(TypeError, match="name` must be a str"):
         backbones(name=float)  # noqa
@@ -59,30 +59,30 @@ def my_model(nc_input=5, nc_output=6):
     assert mlp.weight.shape == (7, 5)
 
     # basic get
-    backbones(my_model, name="cho")
-    assert backbones.get("cho")
+    backbones(my_model, name="foo")
+    assert backbones.get("foo")
 
     # test override
-    backbones(my_model, name="cho", override=True)
-    functions = backbones.get("cho", strict=False)
+    backbones(my_model, name="foo", override=True)
+    functions = backbones.get("foo", strict=False)
     assert len(functions) == 1
 
     # test metadata filtering
-    backbones(my_model, name="cho", namespace="timm", type="resnet")
-    backbones(my_model, name="cho", namespace="torchvision", type="resnet")
-    backbones(my_model, name="cho", namespace="timm", type="densenet")
-    backbones(my_model, name="cho", namespace="timm", type="alexnet")
-    function = backbones.get("cho", with_metadata=True, type="resnet", namespace="timm")
-    assert function["name"] == "cho"
+    backbones(my_model, name="foo", namespace="timm", type="resnet")
+    backbones(my_model, name="foo", namespace="torchvision", type="resnet")
+    backbones(my_model, name="foo", namespace="timm", type="densenet")
+    backbones(my_model, name="foo", namespace="timm", type="alexnet")
+    function = backbones.get("foo", with_metadata=True, type="resnet", namespace="timm")
+    assert function["name"] == "foo"
     assert function["metadata"] == {"namespace": "timm", "type": "resnet"}
 
     # test strict=False and with_metadata=False
-    functions = backbones.get("cho", namespace="timm", strict=False)
+    functions = backbones.get("foo", namespace="timm", strict=False)
     assert len(functions) == 3
     assert all(callable(f) for f in functions)
 
     # test available keys
-    assert backbones.available_keys() == ['cho', 'cho', 'cho', 'cho', 'cho', 'my_model']
+    assert backbones.available_keys() == ['foo', 'foo', 'foo', 'foo', 'foo', 'my_model']
 
 
 # todo (tchaton) Debug this test.

Confusing KerError message for flash registry
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior:

```
from flash.image import ImageClassificationData, ImageClassifier

print(ImageClassifier.backbones.get('abcd'))
```

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

It should throw a keyerror.

### Environment

 - PyTorch Version (e.g., 1.0):
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

### Additional context

Sending in PR.

