import copy
import math
import os

import torch
import torch.nn as nn
import torch.nn.functional as F  # noqa PyPep8Naming

from torch.utils.data import Dataset, DataLoader

CPU_DEVICE = torch.device("cpu")


def normalize_embedding_tensor(embedding):
    return F.normalize(embedding, p=2, dim=1)


def embedding_norm_loss(raw_embedding):
    norms = raw_embedding.norm(dim=1)
    return F.mse_loss(norms, torch.ones_like(norms), reduction="none")


def get_val(x):
    if isinstance(x, torch.Tensor):
        return x.item()
    else:
        return x


def compute_pred_entropy(logits):
    # logits are pre softmax
    p = F.softmax(logits, dim=-1)
    log_p = F.log_softmax(logits, dim=-1)
    return -(p * log_p).sum(dim=-1).mean()


def compute_pred_entropy_clean(logits):
    return float(compute_pred_entropy(logits).item())


def copy_state_dict(state_dict, target_device=None):
    copied_state_dict = copy.deepcopy(state_dict)
    if target_device is None:
        return copied_state_dict
    else:
        # Ensures that tensors with the same data_ptrs point to the same
        # data_ptr after copying
        new_state_dict = {}
        unique_dict = {}
        for k, v in copied_state_dict.items():
            unique_key = tuple(v.shape), v.data_ptr()
            if unique_key not in unique_dict:
                unique_dict[unique_key] = v.to(target_device)
            # Create a view
            new_state_dict[k] = unique_dict[unique_key][:]

        return new_state_dict


def get_parent_child_module_list(model):
    ls = []
    for parent_name, parent_module in model.named_modules():
        for child_name, child_module in parent_module.named_children():
            ls.append((parent_name, parent_module, child_name, child_module))
    return ls


class IdentityModule(nn.Module):
    # noinspection PyMethodMayBeStatic
    def forward(self, *inputs):
        if len(inputs) == 1:
            return inputs[0]
        else:
            return inputs


def set_requires_grad(named_parameters, requires_grad):
    for name, param in named_parameters:
        set_requires_grad_single(param, requires_grad)


def set_requires_grad_single(param, requires_grad):
    param.requires_grad = requires_grad


def get_only_requires_grad(parameters, requires_grad=True):
    if isinstance(parameters, list):
        if not parameters:
            return []
        elif isinstance(parameters[0], tuple):
            return [(n, p) for n, p in parameters if p.requires_grad == requires_grad]
        else:
            return [p for p in parameters if p.requires_grad == requires_grad]
    elif isinstance(parameters, dict):
        return {n: p for n, p in parameters if p.requires_grad == requires_grad}
    else:
        # TODO: Support generators  (issue #1178)
        raise RuntimeError("generators not yet supported")


class ListDataset(Dataset):
    def __init__(self, data: list):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        return self.data[item]


class DataLoaderWithLength(DataLoader):
    def __len__(self):
        # TODO: Revert after https://github.com/pytorch/pytorch/issues/36176 addressed (issue #1179)
        # try:
        #     return super().__len__()
        # except TypeError as e:
        #     try:
        #         return self.get_num_batches()
        #     except TypeError:
        #         pass
        #     raise e
        return self.get_num_batches()

    def get_num_batches(self):
        return math.ceil(len(self.dataset) / self.batch_size)


def is_data_parallel(torch_module):
    return isinstance(torch_module, nn.DataParallel)


def safe_save(obj, path, temp_path=None):
    if temp_path is None:
        temp_path = path + "._temp"
    torch.save(obj, temp_path)
    if os.path.exists(path):
        os.remove(path)
    os.rename(temp_path, path)


def get_model_for_saving(model: nn.Module) -> nn.Module:
    if isinstance(model, nn.DataParallel):
        return model.module
    else:
        return model


def eq_state_dicts(state_dict_1, state_dict_2):
    """Checks if the model weights in state_dict_1 and state_dict_2 are equal.

    Args:
        state_dict_1 (dict): state_dict of a PyTorch model
        state_dict_2 (dict): state_dict of a PyTorch model

    Requires:
        state_dict_1 and state_dict_2 to be from the same model

    Returns:
        bool: Returns True if all model weights are equal in state_dict_1 and state_dict_2
    """
    for key_item_1, key_item_2 in zip(state_dict_1.items(), state_dict_2.items()):
        if torch.equal(key_item_1[1], key_item_2[1]):
            pass
        else:
            return False
    return True
