"""Weight initializer."""
from __future__ import absolute_import, print_function

import re
import logging
import warnings
import json
import numpy as np
from .base import string_types
from .ndarray import NDArray, load
from . import random
from . import registry

# inherit str for backward compatibility
class InitDesc(str):
    """Descriptor for the initialization pattern.

    Parameter
    ---------
    name : str
        Name of variable.
    attrs : dict of str to str
        Attributes of this variable taken from ``Symbol.attr_dict``.
    global_init : Initializer
        Global initializer to fallback to.
    """
    def __new__(cls, name, attrs=None, global_init=None):
        ret = super(InitDesc, cls).__new__(cls, name)
        ret.attrs = attrs or {}
        ret.global_init = global_init
        return ret


class Initializer(object):
    """The base class of an initializer."""
    def __init__(self, **kwargs):
        self._kwargs = kwargs

    def dumps(self):
        """Saves the initializer to string

        Returns
        -------
        str
            JSON formatted string that describes the initializer.

        Examples
        --------
        >>> # Create initializer and retrieve its parameters
        ...
        >>> init = mx.init.Normal(0.5)
        >>> init.dumps()
        '["normal", {"sigma": 0.5}]'
        >>> init = mx.init.Xavier(factor_type="in", magnitude=2.34)
        >>> init.dumps()
        '["xavier", {"rnd_type": "uniform", "magnitude": 2.34, "factor_type": "in"}]'
        """
        return json.dumps([self.__class__.__name__.lower(), self._kwargs])

    def __call__(self, desc, arr):
        """Initialize an array

        Parameters
        ----------
        desc : InitDesc
            Initialization pattern descriptor.

        arr : NDArray
            The array to be initialized.
        """
        if not isinstance(desc, InitDesc):
            self._legacy_init(desc, arr)
            return

        if desc.global_init is None:
            desc.global_init = self
        init = desc.attrs.get('__init__', "")

        if init:
            # when calling Variable initializer
            create(init)._init_weight(desc, arr)
        else:
            # register nnvm::FSetInputVariableAttrs in the backend for new patterns
            # don't add new cases here.
            if desc.endswith('weight'):
                self._init_weight(desc, arr)
            elif desc.endswith('bias'):
                self._init_bias(desc, arr)
            elif desc.endswith('gamma'):
                self._init_gamma(desc, arr)
            elif desc.endswith('beta'):
                self._init_beta(desc, arr)
            else:
                self._init_default(desc, arr)