from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
from tensorflow.keras import layers
import numpy as np


PAD = -1

def scaled_dot_product_attention(q, k, v, mask):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.cast(tf.math.sqrt(dk),matmul_qk.dtype)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += tf.cast((mask * -1e9), scaled_attention_logits.dtype)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights


def point_wise_feed_forward_network(d_model, dff):
    return tf.keras.Sequential([
        layers.Dense(dff, kernel_initializer='glorot_normal', activation='relu'),  # (batch_size, seq_len, dff)
        layers.Dense(d_model, kernel_initializer='glorot_normal')  # (batch_size, seq_len, d_model)
    ])


def get_angles(pos, i, d_model):
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
    return pos * angle_rates


def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encoding = angle_rads[np.newaxis, ...]

    return tf.cast(pos_encoding, dtype=tf.float32)

def create_padding_mask(seq):
    '''
    :param seq: [batch_size * seq_len_k] # k means key in MultiheadAttention
    :return: [batch_size, 1, 1, seq_len_k]
    '''
    if seq.dtype != np.int32:
        print("float")
        seq = tf.cast(tf.math.equal(seq, 0.), tf.float32)
    else:
        print("int")
        seq = tf.cast(tf.math.equal(seq, 0), tf.float32)

    # add extra dimensions so that we can add the padding
    # to the attention logits.
    return seq[:, tf.newaxis,tf.newaxis, :]  # (batch_size, 1,1, seq_len)

def create_look_ahead_mask(size):
    '''
    :param size: == seq_len_k
    :return: (seq_len_q, seq_len_k) 只用在decoderblock1，此时qkv的len全相同，因为block1是对taget的encode
    '''
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)

def create_masks(inp, tar):
    '''
    :param inp: [batch_size * seq_len_k_of_encoder ]
    :param tar: [batch_size * seq_len_q_of_decoder_block1 ]
    :return:
    '''
    # Encoder padding mask
    enc_padding_mask = create_padding_mask(inp)

    # Used in the 2nd attention block in the decoder.
    # This padding mask is used to mask the encoder outputs.
    # encoder outputs [batch_size * seq_len * d_model] 中间那一维相比原始encoder的input不变，所以就按照inp计算了
    dec_padding_mask = create_padding_mask(inp)

    # Used in the 1st attention block in the decoder.
    # It is used to pad and mask future tokens in the input received by the decoder.
    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
    dec_target_padding_mask = create_padding_mask(tar)
    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

    # print('enc_padding_mask',enc_padding_mask)
    # print('combined_mask', combined_mask)
    # print('dec_padding_mask', dec_padding_mask)

    return enc_padding_mask, combined_mask, dec_padding_mask

def create_DecBlock1_pad_mask(tar):
    tar = tf.cast(tf.math.equal(tar, PAD), tf.float32)
    return tar[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1,1, seq_len)

def create_combined_mask(tar):
    # Used in the 1st attention block in the decoder.
    # It is used to pad and mask future tokens in the input received by the decoder.
    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
    dec_target_padding_mask = create_DecBlock1_pad_mask(tar)
    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

    return combined_mask

