from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cntk as C
import numpy as np
from .common import floatx, epsilon, image_dim_ordering, image_data_format
from collections import defaultdict
from contextlib import contextmanager
import warnings


C.set_global_option('align_axis', 1)

b_any = any


dev = C.device.use_default_device()
if dev.type() == 0:
    warnings.warn(
        'CNTK backend warning: GPU is not detected. '
        'CNTK\'s CPU version is not fully optimized,'
        'please run with GPU to get better performance.')

# A learning phase is a bool tensor used to run Keras models in
# either train mode (learning_phase == 1) or test mode (learning_phase == 0).
# LEARNING_PHASE_PLACEHOLDER is the placeholder for dynamic learning phase
_LEARNING_PHASE_PLACEHOLDER = C.constant(shape=(), dtype=np.float32, value=1.0, name='_keras_learning_phase')
# static learning phase flag, if it is not 0 or 1, we will go with dynamic learning phase tensor.
_LEARNING_PHASE = -1
_UID_PREFIXES = defaultdict(int)

# cntk doesn't support gradient as symbolic op, to hook up with keras model,
# we will create gradient as a constant placeholder, here use this global
# map to keep the mapping from grad placeholder to parameter
grad_parameter_dict = {}

NAME_SCOPE_STACK = []


@contextmanager
def name_scope(name):
    global NAME_SCOPE_STACK
    NAME_SCOPE_STACK.append(name)
    yield
    NAME_SCOPE_STACK.pop()


def get_uid(prefix=''):
    _UID_PREFIXES[prefix] += 1
    return _UID_PREFIXES[prefix]


def learning_phase():
    # If _LEARNING_PHASE is not 0 or 1, return dynamic learning phase tensor
    return _LEARNING_PHASE if _LEARNING_PHASE in {0, 1} else _LEARNING_PHASE_PLACEHOLDER


def set_learning_phase(value):
    global _LEARNING_PHASE
    if value not in {0, 1}:
        raise ValueError('CNTK Backend: Set learning phase '
                         'with value %s is not supported, '
                         'expected 0 or 1.' % value)
    _LEARNING_PHASE = value


def clear_session():
    """Reset learning phase flag for cntk backend.
    """
    global _LEARNING_PHASE
    global _LEARNING_PHASE_PLACEHOLDER
    _LEARNING_PHASE = -1
    _LEARNING_PHASE_PLACEHOLDER.value = np.asarray(1.0)


def in_train_phase(x, alt, training=None):
    global _LEARNING_PHASE
    if training is None:
        training = learning_phase()
        uses_learning_phase = True
    else:
        uses_learning_phase = False

    # CNTK currently don't support cond op, so here we use
    # element_select approach as workaround. It may have
    # perf issue, will resolve it later with cntk cond op.
    if callable(x) and isinstance(x, C.cntk_py.Function) is False:
        x = x()
    if callable(alt) and isinstance(alt, C.cntk_py.Function) is False:
        alt = alt()

    if training is True:
        x._uses_learning_phase = uses_learning_phase
        return x
    else:
        # if _LEARNING_PHASE is static
        if isinstance(training, int) or isinstance(training, bool):
            result = x if training == 1 or training is True else alt
        else:
            result = C.element_select(training, x, alt)
        result._uses_learning_phase = uses_learning_phase
        return result


def in_test_phase(x, alt, training=None):
    return in_train_phase(alt, x, training=training)


def _convert_string_dtype(dtype):
    # cntk only support float32 and float64
    if dtype == 'float32':
        return np.float32
    elif dtype == 'float64':
        return np.float64
    else:
        # cntk only running with float,
        # try to cast to float to run the model
        return np.float32


def _convert_dtype_string(dtype):
    if dtype == np.float32:
        return 'float32'
    elif dtype == np.float64:
        return 'float64'
    else:
        raise ValueError('CNTK Backend: Unsupported dtype: %s. '
                         'CNTK only supports float32 and '
                         'float64.' % dtype)


def variable(value, dtype=None, name=None, constraint=None):
    """Instantiates a variable and returns it.

    # Arguments
        value: Numpy array, initial value of the tensor.
        dtype: Tensor type.
        name: Optional name string for the tensor.
        constraint: Optional projection function to be
            applied to the variable after an optimizer update.

    # Returns
        A variable instance (with Keras metadata included).
    """
    if dtype is None:
        dtype = floatx()

    if name is None:
        name = ''

    if isinstance(
            value,
            C.variables.Constant) or isinstance(
            value,
            C.variables.Parameter):
        value = value.value

    # we don't support init parameter with symbolic op, so eval it first as
    # workaround
    if isinstance(value, C.cntk_py.Function):
        value = eval(value)

    shape = value.shape if hasattr(value, 'shape') else ()
    if hasattr(value, 'dtype') and value.dtype != dtype and len(shape) > 0:
        value = value.astype(dtype)

    # TODO: remove the conversion when cntk supports int32, int64
    # https://docs.microsoft.com/en-us/python/api/cntk.variables.parameter
    dtype = 'float32' if 'int' in str(dtype) else dtype

    v = C.parameter(shape=shape,
                    init=value,
                    dtype=dtype,
                    name=_prepare_name(name, 'variable'))
    v._keras_shape = v.shape
    v._uses_learning_phase = False
    v.constraint = constraint
    return v


def bias_add(x, bias, data_format=None):
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {'channels_first', 'channels_last'}:
        raise ValueError('Unknown data_format ' + str(data_format))

    dims = len(x.shape)
    if dims > 0 and x.shape[0] == C.InferredDimension:
        dims -= 1

    bias_dims = len(bias.shape)
    if bias_dims != 1 and bias_dims != dims:
        raise ValueError('Unexpected bias dimensions %d, '
                         'expected 1 or %d dimensions' % (bias_dims, dims))

    if dims == 4:
        if data_format == 'channels_first':
            if bias_dims == 1:
                shape = (bias.shape[0], 1, 1, 1)
            else:
                shape = (bias.shape[3],) + bias.shape[:3]
        elif data_format == 'channels_last':
            if bias_dims == 1:
                shape = (1, 1, 1, bias.shape[0])
            else:
                shape = bias.shape
    elif dims == 3:
        if data_format == 'channels_first':
            if bias_dims == 1:
                shape = (bias.shape[0], 1, 1)
            else:
                shape = (bias.shape[2],) + bias.shape[:2]
        elif data_format == 'channels_last':
            if bias_dims == 1:
                shape = (1, 1, bias.shape[0])
            else:
                shape = bias.shape
    elif dims == 2:
        if data_format == 'channels_first':
            if bias_dims == 1:
                shape = (bias.shape[0], 1)
            else:
                shape = (bias.shape[1],) + bias.shape[:1]
        elif data_format == 'channels_last':
            if bias_dims == 1:
                shape = (1, bias.shape[0])
            else:
                shape = bias.shape
    else:
        shape = bias.shape
    return x + reshape(bias, shape)


def eval(x):
    if isinstance(x, C.cntk_py.Function):
        return x.eval()
    elif isinstance(x, C.variables.Constant) or isinstance(x, C.variables.Parameter):
        return x.value
    else:
        raise ValueError('CNTK Backend: `eval` method on '
                         '`%s` type is not supported. '
                         'CNTK only supports `eval` with '
                         '`Function`, `Constant` or '
                         '`Parameter`.' % type(x))


def placeholder(
        shape=None,
        ndim=None,
        dtype=None,
        sparse=False,
        name=None,
        dynamic_axis_num=1):
    if dtype is None:
        dtype = floatx()
    if not shape:
        if ndim:
            shape = tuple([None for _ in range(ndim)])

    dynamic_dimension = C.FreeDimension if _get_cntk_version() >= 2.2 else C.InferredDimension
    cntk_shape = [dynamic_dimension if s is None else s for s in shape]
    cntk_shape = tuple(cntk_shape)

    if dynamic_axis_num > len(cntk_shape):
        raise ValueError('CNTK backend: creating placeholder with '
                         '%d dimension is not supported, at least '
                         '%d dimensions are needed.'
                         % (len(cntk_shape, dynamic_axis_num)))

    if name is None:
        name = ''

    cntk_shape = cntk_shape[dynamic_axis_num:]

    x = C.input(
        shape=cntk_shape,
        dtype=_convert_string_dtype(dtype),
        is_sparse=sparse,
        name=name)
    x._keras_shape = shape
    x._uses_learning_phase = False
    x._cntk_placeholder = True
    return x


def is_placeholder(x):
    """Returns whether `x` is a placeholder.

    # Arguments
        x: A candidate placeholder.

    # Returns
        Boolean.
    """
    return hasattr(x, '_cntk_placeholder') and x._cntk_placeholder


def is_keras_tensor(x):
    if not is_tensor(x):
        raise ValueError('Unexpectedly found an instance of type `' +
                         str(type(x)) + '`. '
                         'Expected a symbolic tensor instance.')
    return hasattr(x, '_keras_history')


def is_tensor(x):
    return isinstance(x, (C.variables.Constant,
                          C.variables.Variable,
                          C.variables.Parameter,
                          C.ops.functions.Function))


def shape(x):
    shape = list(int_shape(x))
    num_dynamic = _get_dynamic_axis_num(x)
    non_dyn_shape = []
    for i in range(len(x.shape)):
        if shape[i + num_dynamic] is None:
            non_dyn_shape.append(x.shape[i])
        else:
            non_dyn_shape.append(shape[i + num_dynamic])
    return shape[:num_dynamic] + non_dyn_shape


def is_sparse(tensor):
    return tensor.is_sparse


def int_shape(x):
    if hasattr(x, '_keras_shape'):
        return x._keras_shape

    shape = x.shape
    if hasattr(x, 'dynamic_axes'):
        dynamic_shape = [None for a in x.dynamic_axes]
        shape = tuple(dynamic_shape) + shape
    return shape


def ndim(x):
    shape = int_shape(x)
    return len(shape)


def _prepare_name(name, default):
    prefix = '_'.join(NAME_SCOPE_STACK)
    if name is None or name == '':
        return prefix + '/' + default
    return prefix + '/' + name


def constant(value, dtype=None, shape=None, name=None):
    if dtype is None:
        dtype = floatx()
    if shape is None:
        shape = ()
    np_value = value * np.ones(shape)
    const = C.constant(np_value,
                       dtype=dtype,
                       name=_prepare_name(name, 'constant'))
    const._keras_shape = const.shape
    const._uses_learning_phase = False
    return const


def random_binomial(shape, p=0.0, dtype=None, seed=None):
    # use numpy workaround now
    if seed is None:
        # ensure that randomness is conditioned by the Numpy RNG
        seed = np.random.randint(10e7)
        np.random.seed(seed)
    if dtype is None:
        dtype = np.float32
    else:
        dtype = _convert_string_dtype(dtype)

    size = 1
    for _ in shape:
        if _ is None:
            raise ValueError('CNTK Backend: randomness op with '
                             'dynamic shape is not supported now. '
                             'Please provide fixed dimension '
                             'instead of `None`.')
        size *= _

    binomial = np.random.binomial(1, p, size).astype(dtype).reshape(shape)
    return variable(value=binomial, dtype=dtype)


def random_uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
    for _ in shape:
        if _ is None:
            raise ValueError('CNTK Backend: randomness op with '
                             'dynamic shape is not supported now. '
                             'Please provide fixed dimension '
                             'instead of `None`.')

    return random_uniform_variable(shape, minval, maxval, dtype, seed)


def random_uniform_variable(shape, low, high,
                            dtype=None, name=None, seed=None):
    if dtype is None:
        dtype = floatx()
    if seed is None:
        # ensure that randomness is conditioned by the Numpy RNG
        seed = 
np.random.randint(10e3)