# MIT License
#
# Copyright (c) 2023 Transmute AI Lab
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from collections import abc
import torch
from torch._six import string_classes
from torch.nn.parallel.scatter_gather import gather
from torch.nn import DataParallel, Sequential, ModuleList
from torch.nn.parallel import DistributedDataParallel
from trailmet.algorithms.algorithms import BaseAlgorithm

__all__ = ['Distillation']


class Distillation(BaseAlgorithm):

    def __init__(self, **kwargs):
        super(Distillation, self).__init__(**kwargs)

        pass


# https://github.com/yoshitomo-matsubara/torchdistill/blob/8be107e00d6a481711e513f45824ecfaa8249c77/torchdistill/common/module_util.py


def check_if_wrapped(model):
    return isinstance(model, (DataParallel, DistributedDataParallel))


def get_module(root_module, module_path):
    module_names = module_path.split('.')
    module = root_module
    for module_name in module_names:
        if not hasattr(module, module_name):
            if isinstance(module, (DataParallel, DistributedDataParallel)):
                module = module.module
                if not hasattr(module, module_name):
                    if (isinstance(module, Sequential)
                            and module_name.lstrip('-').isnumeric()):
                        module = module[int(module_name)]
                    else:
                        raise Exception(
                            '`{}` of `{}` could not be reached in `{}`'.format(
                                module_name, module_path,
                                type(root_module).__name__))
                        # logger.info('`{}` of `{}` could not be reached in `{}`'.format(module_name, module_path,
                        #                                                                type(root_module).__name__))
                else:
                    module = getattr(module, module_name)
            elif (isinstance(module, (Sequential, ModuleList))
                  and module_name.lstrip('-').isnumeric()):
                module = module[int(module_name)]
            else:
                raise Exception(
                    '`{}` of `{}` could not be reached in `{}`'.format(
                        module_name, module_path,
                        type(root_module).__name__))
                # logger.info('`{}` of `{}` could not be reached in `{}`'.format(module_name, module_path,
                #                                                                type(root_module).__name__))
                return None
        else:
            module = getattr(module, module_name)
    return module


# https://github.com/yoshitomo-matsubara/torchdistill/blob/main/torchdistill/core/forward_hook.py


def get_device_index(data):
    if isinstance(data, torch.Tensor):
        device = data.device
        return 'cpu' if device.type == 'cpu' else device.index
    elif isinstance(data, abc.Mapping):
        for key, data in data.items():
            result = get_device_index(data)
            if result is not None:
                return result
    elif isinstance(data, tuple):
        for d in data:
            result = get_device_index(d)
            if result is not None:
                return result
    elif isinstance(data,
                    abc.Sequence) and not isinstance(data, string_classes):
        for d in data:
            result = get_device_index(d)
            if result is not None:
                return result
    return None


def register_forward_hook_with_dict(module, module_path, requires_input,
                                    requires_output, io_dict):
    io_dict[module_path] = dict()

    def forward_hook4input(self, func_input, func_output):
        if isinstance(func_input, tuple) and len(func_input) == 1:
            func_input = func_input[0]

        device_index = get_device_index(func_output)
        sub_io_dict = io_dict[module_path]
        if 'input' not in sub_io_dict:
            sub_io_dict['input'] = dict()
        sub_io_dict['input'][device_index] = func_input

    def forward_hook4output(self, func_input, func_output):
        if isinstance(func_output, tuple) and len(func_output) == 1:
            func_output = func_output[0]

        device_index = get_device_index(func_output)
        sub_io_dict = io_dict[module_path]
        if 'output' not in sub_io_dict:
            sub_io_dict['output'] = dict()
        sub_io_dict['output'][device_index] = func_output

    def forward_hook4io(self, func_input, func_output):
        if isinstance(func_input, tuple) and len(func_input) == 1:
            func_input = func_input[0]
        if isinstance(func_output, tuple) and len(func_output) == 1:
            func_output = func_output[0]

        device_index = get_device_index(func_output)
        sub_io_dict = io_dict[module_path]
        if 'input' not in sub_io_dict:
            sub_io_dict['input'] = dict()

        if 'output' not in sub_io_dict:
            sub_io_dict['output'] = dict()

        sub_io_dict['input'][device_index] = func_input
        sub_io_dict['output'][device_index] = func_output

    if requires_input and not requires_output:
        return module.register_forward_hook(forward_hook4input)
    elif not requires_input and requires_output:
        return module.register_forward_hook(forward_hook4output)
    elif requires_input and requires_output:
        return module.register_forward_hook(forward_hook4io)
    raise ValueError('Either requires_input or requires_output should be True')


class ForwardHookManager(object):
    """
    Example::
        >>> import torch
        >>> from torchvision import models
        >>> from torchdistill.core.forward_hook import ForwardHookManager
        >>> device = torch.device('cpu')
        >>> forward_hook_manager = ForwardHookManager(device)
        >>> model = models.resnet18()
        >>> forward_hook_manager.add_hook(model, 'layer2')
        >>> x = torch.rand(16, 3, 224, 224)
        >>> y = model(x)
        >>> io_dict = forward_hook_manager.pop_io_dict()
        >>> layer2_input_tensor = io_dict['layer2']['input']
        >>> layer2_output_tensor = io_dict['layer2']['output']
    """

    def __init__(self, target_device):
        self.target_device = (torch.device(target_device) if isinstance(
            target_device, str) else target_device)
        self.uses_cuda = self.target_device.type == 'cuda'
        self.io_dict = dict()
        self.hook_list = list()

    def add_hook(self,
                 module,
                 module_path,
                 requires_input=True,
                 requires_output=True):
        unwrapped_module = module.module if check_if_wrapped(
            module) else module
        sub_module = get_module(unwrapped_module, module_path)
        handle = register_forward_hook_with_dict(sub_module, module_path,
                                                 requires_input,
                                                 requires_output, self.io_dict)
        self.hook_list.append((module_path, handle))

    def pop_io_dict(self):
        gathered_io_dict = dict()
        for module_path, module_io_dict in self.io_dict.items():
            gathered_io_dict[module_path] = dict()
            for io_type in list(module_io_dict.keys()):
                sub_dict = module_io_dict.pop(io_type)
                values = [sub_dict[key] for key in sorted(sub_dict.keys())]
                gathered_obj = (gather(values, self.target_device)
                                if self.uses_cuda and len(values) > 1 else
                                values[-1])
                gathered_io_dict[module_path][io_type] = gathered_obj
        return gathered_io_dict

    def pop_io_dict_from_device(self, device):
        device_io_dict = dict()
        device_key = device.index if device.type == 'cuda' else device.type
        for module_path, module_io_dict in self.io_dict.items():
            device_io_dict[module_path] = dict()
            for io_type in list(module_io_dict.keys()):
                sub_dict = module_io_dict[io_type]
                device_io_dict[module_path][io_type] = sub_dict.pop(device_key)
        return device_io_dict

    def change_target_device(self, target_device):
        if self.target_device.type != target_device.type:
            for sub_dict in self.io_dict.values():
                sub_dict.clear()
        self.target_device = target_device

    def clear(self):
        self.io_dict.clear()
        for _, handle in self.hook_list:
            handle.remove()
        self.hook_list.clear()
