# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# The 1.6 release of PyTorch switched torch.save to use a new zipfile-based
# file format. It will cause RuntimeError when a checkpoint was saved in
# torch >= 1.6.0 but loaded in torch < 1.7.0.
# More details at https://github.com/open-mmlab/mmpose/issues/904

from ..path import mkdir_or_exist
from ..version_utils import digit_version
from .parrots_wrapper import TORCH_VERSION

if TORCH_VERSION != 'parrots' and digit_version(TORCH_VERSION) < digit_version(
        '1.7.0'):
    # Modified from https://github.com/pytorch/pytorch/blob/master/torch/hub.py
    import os
    import sys
    import warnings
    import zipfile
    from urllib.parse import urlparse

    import torch
    from torch.hub import HASH_REGEX, _get_torch_home, download_url_to_file

    # Hub used to support automatically extracts from zipfile manually
    # compressed by users. The legacy zip format expects only one file from
    # torch.save() < 1.6 in the zip. We should remove this support since
    # zipfile is now default zipfile format for torch.save().
    def _is_legacy_zip_format(filename):
        if zipfile.is_zipfile(filename):
            infolist = zipfile.ZipFile(filename).infolist()
            return len(infolist) == 1 and not infolist[0].is_dir()
        return False

    def _legacy_zip_load(filename, model_dir, map_location):
        warnings.warn(
            'Falling back to the old format < 1.6. This support will'
            ' be deprecated in favor of default zipfile format '
            'introduced in 1.6. Please redo torch.save() to save it '
            'in the new zipfile format.', DeprecationWarning)
        # Note: extractall() defaults to overwrite file if exists. No need to
        #       clean up beforehand. We deliberately don't handle tarfile here
        #       since our legacy serialization format was in tar.
        #       E.g. resnet18-5c106cde.pth which is widely used.
        with zipfile.ZipFile(filename) as f:
            members = f.infolist()
            if len(members) != 1:
                raise RuntimeError(
                    'Only one file(not dir) is allowed in the zipfile')
            f.extractall(model_dir)
            extraced_name = members[0].filename
            extracted_file = os.path.join(model_dir, extraced_name)
        return torch.load(extracted_file, map_location=map_location)

    def load_url(url,
                 model_dir=None,
                 map_location=None,
                 progress=True,
                 check_hash=False,
                 file_name=None):
        r"""Loads the Torch serialized object at the given URL.
        If downloaded file is a zip file, it will be automatically decompressed
        If the object is already present in `model_dir`, it's deserialized and
        returned.
        The default value of ``model_dir`` is ``<hub_dir>/checkpoints`` where
        ``hub_dir`` is the directory returned by :func:`~torch.hub.get_dir`.
        Args:
            url (str): URL of the object to download
            model_dir (str, optional): directory in which to save the object
            map_location (optional): a function or a dict specifying how to
                remap storage locations (see torch.load)
            progress (bool, optional): whether or not to display a progress bar
                to stderr. Defaults to True
            check_hash(bool, optional): If True, the filename part of the URL
                should follow the naming convention ``filename-<sha256>.ext``
                where ``<sha256>`` is the first eight or more digits of the
                SHA256 hash of the contents of the file. The hash is used to
                ensure unique names and to verify the contents of the file.
                Defaults to False
            file_name (str, optional): name for the downloaded file. Filename
                from ``url`` will be used if not set. Defaults to None.
        Example:
            >>> url = ('https://s3.amazonaws.com/pytorch/models/resnet18-5c106'
            ...        'cde.pth')
            >>> state_dict = torch.hub.load_state_dict_from_url(url)
        """
        # Issue warning to move data if old env is set
        if os.getenv('TORCH_MODEL_ZOO'):
            warnings.warn(
                'TORCH_MODEL_ZOO is deprecated, please use env '
                'TORCH_HOME instead', DeprecationWarning)

        if model_dir is None:
            torch_home = _get_torch_home()
            model_dir = os.path.join(torch_home, 'checkpoints')

        mkdir_or_exist(model_dir)

        parts = urlparse(url)
        filename = os.path.basename(parts.path)
        if file_name is not None:
            filename = file_name
        cached_file = os.path.join(model_dir, filename)
        if not os.path.exists(cached_file):
            sys.stderr.write('Downloading: "{}" to {}\n'.format(
                url, cached_file))
            hash_prefix = None
            if check_hash:
                r = HASH_REGEX.search(filename)  # r is Optional[Match[str]]
                hash_prefix = r.group(1) if r else None
            download_url_to_file(
                url, cached_file, hash_prefix, progress=progress)

        if _is_legacy_zip_format(cached_file):
            return _legacy_zip_load(cached_file, model_dir, map_location)

        try:
            return torch.load(cached_file, map_location=map_location)
        except RuntimeError as error:
            if digit_version(TORCH_VERSION) < digit_version('1.5.0'):
                warnings.warn(
                    f'If the error is the same as "{cached_file} is a zip '
                    'archive (did you mean to use torch.jit.load()?)", you can'
                    ' upgrade your torch to 1.5.0 or higher (current torch '
                    f'version is {TORCH_VERSION}). The error was raised '
                    ' because the checkpoint was saved in torch>=1.6.0 but '
                    'loaded in torch<1.5.')
            raise error
else:
    from torch.utils.model_zoo import load_url  # type: ignore # noqa: F401