import math
import numpy as np
from scipy.stats import multivariate_normal
import tensorflow as tf

def learning_rate(step_num, d_model=512, warmup_steps=4000):
    a = step_num**(-0.5)
    b = step_num*warmup_steps**(-1.5)
    return a, b, d_model**(-0.5) * min(step_num**(-0.5), step_num*(warmup_steps**(-1.5))) 

def selu(x):
    alpha = 1.6732632423543772848170429916717
    scale = 1.0507009873554804934193349852946
    print('use selu')
    return scale*tf.where(x>=0.0, x, alpha*tf.nn.elu(x))

def bilinear_sim_4d(x, y, is_nor=True):
    '''calulate bilinear similarity with two 4d tensor.
    
    Args:
        x: a tensor with shape [batch, time_x, dimension_x, num_stacks]
        y: a tensor with shape [batch, time_y, dimension_y, num_stacks]

    Returns:
        a tensor with shape [batch, time_x, time_y, num_stacks]

    Raises:
        ValueError: if
            the shapes of x and y are not match;
            bilinear matrix reuse error.
    '''
    M = tf.get_variable(
        name="bilinear_matrix", 
        shape=[x.shape[2], y.shape[2], x.shape[3]],
        dtype=tf.float32,
        initializer=tf.orthogonal_initializer())
    sim = tf.einsum('biks,kls,bjls->bijs', x, M, y)

    if is_nor:
        scale = tf.sqrt(tf.cast(x.shape[2] * y.shape[2], tf.float32))
        scale = tf.maximum(1.0, scale)
        return sim / scale
    else:
        return sim


def bilinear_sim(x, y, is_nor=True):
    '''calculate bilinear similarity with two tensor.
    Args:
        x: a tensor with shape [batch, time_x, dimension_x]
        y: a tensor with shape [batch, time_y, dimension_y]
    
    Returns:
        a tensor with shape [batch, time_x, time_y]
    Raises:
        ValueError: if
            the shapes of x and y are not match;
            bilinear matrix reuse error.
    '''
    M = tf.get_variable(
        name="bilinear_matrix", 
        shape=[x.shape[-1], y.shape[-1]],
        dtype=tf.float32,
        initializer=tf.orthogonal_initializer())
    sim = tf.einsum('bik,kl,bjl->bij', x, M, y)

    if is_nor:
        scale = tf.sqrt(tf.cast(x.shape[-1] * y.shape[-1], tf.float32))
        scale = tf.maximum(1.0, scale)
        return sim / scale
    else:
        return sim

def dot_sim(x, y, is_nor=True):
    '''calculate dot similarity with two tensor.

    Args:
        x: a tensor with shape [batch, time_x, dimension]
        y: a tensor with shape [batch, time_y, dimension]
    
    Returns:
        a tensor with shape [batch, time_x, time_y]
    Raises:
        AssertionError: if
            the shapes of x and y are not match.
    '''
    assert x.shape[-1] == y.shape[-1]

    sim = tf.einsum('bik,bjk->bij', x, y)

    if is_nor:
        scale = tf.sqrt(tf.cast(x.shape[-1], tf.float32))
        scale = tf.maximum(1.0, scale)
        return sim / scale
    else:
        return sim

def layer_norm(x, axis=None, epsilon=1e-6):
    '''Add layer normalization.

    Args:
        x: a tensor
        axis: the dimensions to normalize

    Returns:
        a tensor the same shape as x.

    Raises:
    '''
    print('wrong version of layer_norm')
    scale = tf.get_variable(
        name='scale',
        shape=[1],
        dtype=tf.float32,
        initializer=tf.ones_initializer())
    bias = tf.get_variable(
        name='bias',
        shape=[1],
        dtype=tf.float32,
        initializer=tf.zeros_initializer())

    if axis is None:
        axis = [-1]

    mean = tf.reduce_mean(x, axis=axis, keep_dims=True)
    variance = tf.reduce_mean(tf.square(x - mean), axis=axis, keep_dims=True)
    norm = (x-mean) * tf.rsqrt(variance + epsilon)
    return scale * norm + bias

def layer_norm_debug(x, axis = None, epsilon=1e-6):
    '''Add layer normalization.

    Args:
        x: a tensor
        axis: the dimensions to normalize

    Returns:
        a tensor the same shape as x.

    Raises:
    '''
    if axis is None:
        axis = [-1]
    shape = [x.shape[i] for i in axis]

    scale = tf.get_variable(
        name='scale',
        shape=shape,
        dtype=tf.float32,
        initializer=tf.ones_initializer())
    bias = tf.get_variable(
        name='bias',
        shape=shape,
        dtype=tf.float32,
        initializer=tf.zeros_initializer())

    mean = tf.reduce_mean(x, axis=axis, keep_dims=True)
    variance = tf.reduce_mean(tf.square(x - mean), axis=axis, keep_dims=True)
    norm = (x-mean) * tf.rsqrt(variance + epsilon)
    return scale * norm + bias

def dense(x, out_dimension=None, add_bias=True):
    '''Add dense connected layer, Wx + b.

    Args:
        x: a tensor with shape [batch, time, dimension]
        out_dimension: a number which is the output dimension

    Return:
        a tensor with shape [batch, time, out_dimension]

    Raises:
    '''
    if out_dimension is None:
        out_dimension = x.shape[-1]

    W = tf.get_variable(
        name='weights',
        shape=[x.shape[-1], out_dimension],
        dtype=tf.float32,
        # initializer=tf.orthogonal_initializer())
        initializer=tf.contrib.layers.xavier_initializer())
    if add_bias:
        bias = tf.get_variable(
            name='bias',
            shape=[1],
            dtype=tf.float32,
            initializer=tf.zeros_initializer())
        return tf.einsum('bik,kj->bij', x, W) + bias
    else:
        return tf.einsum('bik,kj->bij', x, W)

def matmul_2d(x, out_dimension, drop_prob=None):
    '''Multiplies 2-d tensor by weights.

    Args:
        x: a tensor with shape [batch, dimension]
        out_dimension: a number

    Returns:
        a tensor with shape [batch, out_dimension]

    Raises:
    '''
    W = tf.get_variable(
        name='weights',
        shape=[x.shape[1], out_dimension],
        dtype=tf.float32,
        initializer=tf.orthogonal_initializer())
    if drop_prob is not None:
        W = tf.nn.dropout(W, drop_prob)
        print('W is dropout')

    return tf.matmul(x, W)

def gauss_positional_encoding_vector(x, role=0, value=0):
    position = int(x.shape[1])
    dimension = int(x.shape[2])
    print('position: %s' %position)
    print('dimension: %s' %dimension)

    _lambda = tf.get_variable(
        name='lambda',
        shape=[position],
        dtype=tf.float32,
        initializer=tf.constant_initializer(value))
    _lambda = tf.expand_dims(_lambda, axis=-1)

    mean = [position/2.0, dimension/2.0]

    #cov = [[position/3.0, 0], [0, dimension/3.0]]
    sigma_x = position/math.sqrt(4.0*dimension)
    sigma_y = math.sqrt(dimension/4.0)
    cov = [[sigma_x*sigma_x, role*sigma_x*sigma_y], 
            [role*sigma_x*sigma_y, sigma_y*sigma_y]]

    pos = np.dstack(np.mgrid[0:position, 0:dimension])

    
    rv = multivariate_normal(mean, cov)
    signal = rv.pdf(pos) 
    signal = signal - np.max(signal)/2.0

    signal = tf.multiply(_lambda, signal)
    signal = tf.expand_dims(signal, axis=0)

    print('gauss positional encoding')

    return x + _lambda * signal

def positional_encoding(x, min_timescale=1.0, max_timescale=1.0e4, value=0):
    '''Adds a bunch of sinusoids of different frequencies to a tensor.

    Args:
        x: a tensor with shape [batch, length, channels]
        min_timescale: a float
        max_timescale: a float

    Returns:
        a tensor the same shape as x.

    Raises:
    '''
    length = x.shape[1]
    channels = x.shape[2]
    _lambda = tf.get_variable(
        name='lambda',
        shape=[1],
        dtype=tf.float32,
        initializer=tf.constant_initializer(value))

    position = tf.to_float(tf.range(length))
    num_timescales = channels // 2
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.to_float(num_timescales) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])
    #signal = tf.reshape(signal, [1, length, channels])
    signal = tf.expand_dims(signal, axis=0)

    return x + _lambda * signal


def positional_encoding_vector(x, min_timescale=1.0, max_timescale=1.0e4, value=0):
    '''Adds a bunch of sinusoids of different frequencies to a tensor.

    Args:
        x: a tensor with shape [batch, length, channels]
        min_timescale: a float
        max_timescale: a float

    Returns:
        a tensor the same shape as x.

    Raises:
    '''
    length = x.shape[1]
    channels = x.shape[2]
    _lambda = tf.get_variable(
        name='lambda',
        shape=[length],
        dtype=tf.float32,
        initializer=tf.constant_initializer(value))
    _lambda = tf.expand_dims(_lambda, axis=-1)

    position = tf.to_float(tf.range(length))
    num_timescales = channels // 2
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.to_float(num_timescales) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]])

    signal = tf.multiply(_lambda, signal)
    signal = tf.expand_dims(signal, axis=0)

    return x + signal

def mask(row_lengths, col_lengths, max_row_length, max_col_length):
    '''Return a mask tensor representing the first N positions of each row and each column.

    Args:
        row_lengths: a tensor with shape [batch]
        col_lengths: a tensor with shape [batch]

    Returns:
        a mask tensor with shape [batch, max_row_length, max_col_length]

    Raises:
    '''
    row_mask = tf.sequence_mask(row_lengths, max_row_length) #bool, [batch, max_row_len]
    col_mask = tf.sequence_mask(col_lengths, max_col_length) #bool, [batch, max_col_len]

    row_mask = tf.cast(tf.expand_dims(row_mask, -1), tf.float32)
    col_mask = tf.cast(tf.expand_dims(col_mask, -1), tf.float32)

    return tf.einsum('bik,bjk->bij', row_mask, col_mask)

def weighted_sum(weight, values):
    '''Calcualte the weighted sum.

    Args:
        weight: a tensor with shape [batch, time, dimension]
        values: a tensor with shape [batch, dimension, values_dimension]

    Return:
        a tensor with shape [batch, time, values_dimension]

    Raises:
    '''
    return tf.einsum('bij,bjk->bik', weight, values)


        

