from typing import Union, List
import torch


def is_differentiable(
        loss: torch.Tensor, model: Union[torch.nn.Module, List[torch.nn.Module]], print_instead: bool = False
) -> None:
    """
    Overview:
        Judge whether the model/models are differentiable. First check whether module's grad is None,
        then do loss's back propagation, finally check whether module's grad are torch.Tensor.
    Arguments:
        - loss (:obj:`torch.Tensor`): loss tensor of the model
        - model (:obj:`Union[torch.nn.Module, List[torch.nn.Module]]`): model or models to be checked
        - print_instead (:obj:`bool`): Whether to print module's final grad result, \
            instead of asserting. Default set to ``False``.
    """
    assert isinstance(loss, torch.Tensor)
    if isinstance(model, list):
        for m in model:
            assert isinstance(m, torch.nn.Module)
            for k, p in m.named_parameters():
                assert p.grad is None, k
    elif isinstance(model, torch.nn.Module):
        for k, p in model.named_parameters():
            assert p.grad is None, k
    else:
        raise TypeError('model must be list or nn.Module')

    loss.backward()

    if isinstance(model, list):
        for m in model:
            for k, p in m.named_parameters():
                if print_instead:
                    if not isinstance(p.grad, torch.Tensor):
                        print(k, "grad is:", p.grad)
                else:
                    assert isinstance(p.grad, torch.Tensor), k
    elif isinstance(model, torch.nn.Module):
        for k, p in model.named_parameters():
            if print_instead:
                if not isinstance(p.grad, torch.Tensor):
                    print(k, "grad is:", p.grad)
            else:
                assert isinstance(p.grad, torch.Tensor), k