# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.

import codecs
import logging

import torch
import torch.nn as nn
import torch.nn.functional as F

from abc import ABC, abstractmethod

from utils.exceptions import LayerDefineError, ConfigurationError


class BaseConf(ABC):
    """Basic configuration

    Args:
        input_dim (int): the dimension of input.
        hidden_dim (int): the dimension of hidden state.
        dropout (float): dropout rate.
        (others)...
    """
    def __init__(self, **kwargs):
        self.default()

        for key in kwargs:
            setattr(self, key, kwargs[key])

        self.declare()

        self.verify_before_inference()

        # Moved to get_conf() in Model.py
        #self.inference()
        #self.verify()

    def default(self):
        """ Define the default hyper parameters here. You can define these hyper parameters in your configuration file as well.

        Returns:
            None

        """
        #self.input_dims = [xxx, xxx]  would be inferenced automatically
        self.hidden_dim = 10


    @abstractmethod
    def declare(self):
        """ Define things like "input_ranks" and "num_of_inputs", which are certain with regard to your layer

            num_of_input is N(N>0) means this layer accepts N inputs;

            num_of_input is -1 means this layer accepts any number of inputs;

            The rank here is not the same as matrix rank:

              For a scalar, its rank is 0;\n
              For a vector, its rank is 1;\n
              For a matrix, its rank is 2;\n
              For a cube of numbers, its rank is 3.\n
            ...
            For instance, the rank of (batch size, sequence length, hidden_dim) is 3.

            if num_of_input > 0:

              len(input_ranks) should be equal to num_of_input

            elif num_of_input == -1:

              input_ranks should be a list with only one element and the rank of all the inputs should be equal to that element.

            NOTE: when we build the model, if num_of_input is -1, we would replace it with the real number of inputs and replace input_ranks with a list of real input_ranks.

        Returns:
            None

        """
        self.num_of_inputs = 1
        self.input_ranks = [3]

    @abstractmethod
    def inference(self):
        """ Inference things like output_dim, which may relies on defined hyper parameter such as hidden dim and input_dim

        Returns:
            None

        """
        # inference the output dim and output rank from inputs. Here are some examples:
        #self.output_dim = copy.deepcopy(self.input_dims[0])
        #self.output_dim[-1] = sum([input_dim[-1] for input_dim in self.input_dims])

        self.output_rank = len(self.output_dim)  # DON'T MODIFY THIS

    def verify_before_inference(self):
        """ Some conditions must be fulfilled, otherwise there would be errors when calling inference()

        The difference between verify_before_inference() and verify() is that:
            verify_before_inference() is called before inference() while verify() is called after inference().

        Returns:
            None

        """
        necessary_attrs_for_dev = ['num_of_inputs', 'input_ranks']
        for attr in necessary_attrs_for_dev:
            self.add_attr_exist_assertion_for_dev(attr)

        type_checks = [('num_of_inputs', int),
                       ('input_ranks', list)]
        for attr, attr_type in type_checks:
            self.add_attr_type_assertion(attr, attr_type)


    def verify(self):
        """ Define some necessary varification for your layer when we define the model.

        If you define your own layer and rewrite this funciton, please add "super(YourLayerConf, self).verify()" at the beginning

        Returns:
            None

        """
        self.verify_before_inference()      # prevent to be neglected, verify again

        necessary_attrs_for_user = []
        for attr in necessary_attrs_for_user:
            self.add_attr_exist_assertion_for_user(attr)

        necessary_attrs_for_dev = ['input_dims', 'output_dim', 'output_rank', 'use_gpu']
        for attr in necessary_attrs_for_dev:
            self.add_attr_exist_assertion_for_dev(attr)

        type_checks = [('output_dim', list),
                       ('input_dims', list),
                       ('output_rank', int)]
        for attr, attr_type in type_checks:
            self.add_attr_type_assertion(attr, attr_type)

        range_checks = [('dropout', (0, 1), (True, True))]
        for attr, ranges, bound_legal in range_checks:
            self.add_attr_range_assertion(attr, ranges, bound_legal)

        # demonstration for value checks
        '''
        value_checks = [('some_attr', ['legal_value1', 'legal_value2', 'legal_value3'])]
        for attr, legal_values in value_checks:
            self.add_attr_value_assertion(attr, legal_values)
        '''

        # To check if deepcopy is applied
        assert id(self.output_dim) != id(self.input_dims[0]), 'Please use copy.deepcopy to copy the input_dim to output_dim'

    def verify_former_block(self, former_conf):
        """check if has special attribute rely on former layer

        """
        return True

    def add_attr_type_assertion(self, attr, specified_type):
        """ check if the types of attributes are legal

        Args:
            attr (str): the attribution name
            specified_type (None/str/list): one specified_type of a list of specified_type(including None)

        Returns:
            None

        """
        if not hasattr(self, attr):
            return

        if isinstance(specified_type, list):
            hit_flag = False
            for specified_type_ in specified_type:
                if specified_type_ is None:
                    if getattr(self, attr) is None:
                        hit_flag = True
                        break
                else:
                    if isinstance(getattr(self, attr), specified_type_):
                        hit_flag = True
                        break

            if hit_flag is False:
                raise Exception("For layer %s, the attribute %s should be one of [%s]!" % (
                    type(self).__name__, attr, ", ".join(specified_type_.__name__ if specified_type_ is not None else "None" for specified_type_ in specified_type)))
        else:
            if not (getattr(self, attr) is None and specified_type is None or isinstance(getattr(self, attr), specified_type)):
                raise LayerDefineError("For layer %s, the attribute %s should be a/an %s!" %
                            (type(self).__name__, attr, specified_type.__name__))

    def add_attr_range_assertion(self, attr, range, bounds_legal=(True, True)):
        """ check if attribute falls into the legal range

        Args:
            attr (str): the attribution name
            range (tuple): (num/float('-inf')/None, num/float('inf')/None), None means -inf or inf.
            bounds_legal (tuple): (bool, bool), if the left/right bound is legal

        Returns:
            None

        """
        if not hasattr(self, attr):
            return

        value = getattr(self, attr)
        range = list(range)
        bounds_legal = list(bounds_legal)
        if range[0] is None:
            range[0] = float('-inf')
        if range[1] is None:
            range[1] = float('inf')
        if range[0] == float('-inf'):
            bounds_legal[0] = False
        if range[1] == float('inf'):
            bounds_legal[1] = False

        left_bound_ch = '[' if bounds_legal[0] else '('
        right_bound_ch = ']' if bounds_legal[1] else ')'
        if not ((bounds_legal[0] and value >= range[0] or bounds_legal[0] is False and value > range[0]) and (
                bounds_legal[1] and value <= range[1] or bounds_legal[1] is False and value < range[1])):
            raise Exception("For layer %s, the legal range of attribute %s is %s%f, %f%s" % (
                type(self).__name__, attr, left_bound_ch, range[0], range[1], right_bound_ch))

    def add_attr_exist_assertion_for_dev(self, attr):
        """ check if there are some attributes being forgot by developers

        Args:
            attr (str): the attribution name

        Returns:
            None

        """
        if not hasattr(self, attr):
            raise LayerDefineError("For layer %s, please define %s attribute in declare() or inference()!" % (type(self).__name__, attr))

    def add_attr_exist_assertion_for_user(self, attr):
        """ check if there are some attributes being forgot by users

        Args:
            attr (str): the attribution name

        Returns:
            None

        """
        if not hasattr(self, attr):
            raise ConfigurationError("For layer %s, please configure %s attribute for %s in the configuration file!" % (type(self).__name__, attr, type(self).__name__))

    def add_attr_value_assertion(self, attr, legal_values):
        """ check if attr equals to one of the legal values

        Args:
            attr (str): the attribution name
            legal_values (list): include the legal value

        Returns:
            None

        """
        if not hasattr(self, attr):
            return
        hit_flag = False
        for legal_value in legal_values:
            if getattr(self, attr) == legal_value:
                hit_flag = True
                break
        if hit_flag is False:
            raise Exception("For layer %s, attribute %s should be one of [%s], but you give %s." % \
                (type(self).__name__, attr, ", ".join(str(legal_value) for legal_value in legal_values), str(getattr(self, attr))))


class BaseLayer(nn.Module):
    """The base class of layers

    Args:
        layer_conf (BaseConf): configuration of a layer
    """
    def __init__(self, layer_conf):
        super(BaseLayer, self).__init__()
        self.layer_conf = layer_conf

    def forward(self, *args):
        """

        Args:
            *args (list): a list of args in which arg should be a pair of (representation, length)

        Returns:
            None

        """
        pass

    def is_cuda(self):
        """ To judge if the layer is on CUDA
        if there are parameters in this layer, judge according to the parameters;
        else: judge according to the self.layer_conf.use_gpu

        Returns:
            bool: whether to use gpu

        """
        try:
            # In case someone forget to check use_gpu flag, this function would first detect if the parameters are on cuda
            # ret = next(self.parameters()).data.is_cuda
            ret = self.layer_conf.use_gpu
        except StopIteration as e:
            if not hasattr(self, 'layer_conf'):
                logging.error('Layer.layer_conf must be defined!')
            else:
                logging.error(e)

        return ret

