from collections import abc
import numpy as np
import torch
import functools
from inspect import getfullargspec



def cast_tensor_type(inputs, src_type, dst_type):
    if isinstance(inputs, torch.Tensor):
        return inputs.to(dst_type)
    elif isinstance(inputs, str):
        return inputs
    elif isinstance(inputs, np.ndarray):
        return inputs
    elif isinstance(inputs, abc.Mapping):
        return type(inputs)({
            k: cast_tensor_type(v, src_type, dst_type)
            for k, v in inputs.items()
        })
    elif isinstance(inputs, abc.Iterable):
        return type(inputs)(
            cast_tensor_type(item, src_type, dst_type) for item in inputs)
    else:
        return inputs


def auto_fp16(apply_to=None, out_fp32=False):
  """Decorator to enable fp16 training automatically.

  This decorator is useful when you write custom modules and want to support
  mixed precision training. If inputs arguments are fp32 tensors, they will
  be converted to fp16 automatically. Arguments other than fp32 tensors are
  ignored.

  Args:
      apply_to (Iterable, optional): The argument names to be converted.
          `None` indicates all arguments.
      out_fp32 (bool): Whether to convert the output back to fp32.

  :Example:

      class MyModule1(nn.Module)

          # Convert x and y to fp16
          @auto_fp16()
          def forward(self, x, y):
              pass

      class MyModule2(nn.Module):

          # convert pred to fp16
          @auto_fp16(apply_to=('pred', ))
          def do_something(self, pred, others):
              pass
  """

  def auto_fp16_wrapper(old_func):

    @functools.wraps(old_func)
    def new_func(*args, **kwargs):
      # check if the module has set the attribute `fp16_enabled`, if not,
      # just fallback to the original method.
      if not isinstance(args[0], torch.nn.Module):
        raise TypeError('@auto_fp16 can only be used to decorate the '
                        'method of nn.Module')
      if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled):
        return old_func(*args, **kwargs)
      # get the arg spec of the decorated method
      args_info = getfullargspec(old_func)
      # get the argument names to be casted
      args_to_cast = args_info.args if apply_to is None else apply_to
      # convert the args that need to be processed
      new_args = []
      # NOTE: default args are not taken into consideration
      if args:
        arg_names = args_info.args[:len(args)]
        for i, arg_name in enumerate(arg_names):
          if arg_name in args_to_cast:
            new_args.append(
              cast_tensor_type(args[i], torch.float, torch.half))
          else:
            new_args.append(args[i])
      # convert the kwargs that need to be processed
      new_kwargs = {}
      if kwargs:
        for arg_name, arg_value in kwargs.items():
          if arg_name in args_to_cast:
            new_kwargs[arg_name] = cast_tensor_type(
              arg_value, torch.float, torch.half)
          else:
            new_kwargs[arg_name] = arg_value
      # apply converted arguments to the decorated method
      output = old_func(*new_args, **new_kwargs)
      # cast the results back to fp32 if necessary
      if out_fp32:
        output = cast_tensor_type(output, torch.half, torch.float)
      return output

    return new_func

  return auto_fp16_wrapper





