"""
# **Main ideas is as follows：**

1.Basic transformer：  Portugese-English translation . I consider contend_id as Portugese，answered_correctly as English, so SAINT is contend_id -answered_correctly translation.


2.modify create_masks function

3.Trained on 12M data.

"""
"""
The reference of this notebook is as follows :

1.https://tensorflow.google.cn/tutorials/text/transformer

2.our host's papers (https://arxiv.org/abs/2002.07033, https://arxiv.org/abs/2010.12042)
"""
"""
# **SANIT++**
1.Encoder add features:

part\tags1\tags2

user_lecture_lv

2 Decoder add features:

lagtime\lagtime2\lagtime3

elapsed time\had explation

"""
!pip install ../input/python-datatable/datatable-0.11.0-cp37-cp37m-manylinux2010_x86_64.whl > /dev/null 2>&1
import datatable as dt
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import tensorflow as tf
from sklearn.metrics import roc_auc_score
from collections import defaultdict
import gc

import time
import matplotlib.pyplot as plt

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
"""
# **Data**
"""
content_count=13523
target_count=2
user_lecture_count=500
lagtime_count=20002
elapsed_time_count=500

MAX_LENGTH=50

BUFFER_SIZE = 2000
BATCH_SIZE = 128

#train parameter
#The values used in the base model of transformer were:
#num_layers=6, d_model = 512, dff = 2048. See the paper for all the other versions of the transformer.

#SAINT:The window size, dropout rate, and batch size are set to 100, 0.1, and 128 respectively.
#The best performing model of SAINT has 4 layers and a latent space dimension of 512.
num_layers = 2
d_model = 128
dff = 512
num_heads = 8

# input_vocab_size = content_count + 1
# target_vocab_size = target_count + 1

#content_id have 0，so add 1
input_vocab_size = content_count+1
target_vocab_size = target_count+1
dropout_rate = 0.1

#
EPOCHS = 10
train_flag=False
target = 'answered_correctly'
%%time
data_types_dict = {
    'timestamp': 'int64',
    'user_id': 'int32', 
    'content_id': 'int16', 
    'content_type_id':'int8', 
    #'task_container_id': 'int16',
    #'user_answer': 'int8',
    'answered_correctly': 'int8', 
    'prior_question_elapsed_time': 'float32', 
    'prior_question_had_explanation': 'bool'
}

print('start read train data...')
train_df = dt.fread('../input/riiid-test-answer-prediction/train.csv', columns=set(data_types_dict.keys())).to_pandas()
user_lecture_agg = train_df.groupby('user_id')['content_type_id'].agg(['sum', 'count'])
user_lecture_agg=user_lecture_agg.astype('int16')

user_lecture_sum_dict = user_lecture_agg['sum'].astype('int16').to_dict(defaultdict(int))
user_lecture_count_dict = user_lecture_agg['count'].astype('int16').to_dict(defaultdict(int))

del user_lecture_agg
cum = train_df.groupby('user_id')['content_type_id'].agg(['cumsum', 'cumcount'])
cum['cumcount']=cum['cumcount']+1

#train_df['user_lecture_lv'] = cum['cumsum'] / cum['cumcount']
train_df['user_lecture_lv'] = cum['cumsum'] 
train_df.user_lecture_lv=train_df.user_lecture_lv.astype('int16')
del cum
train_df['user_lecture_lv']
train_df['prior_question_had_explanation'].fillna(2, inplace=True)
train_df = train_df.astype(data_types_dict)
train_df = train_df[train_df[target] != -1].reset_index(drop=True)
max_timestamp_u = train_df[['user_id','timestamp']].groupby(['user_id']).agg(['max']).reset_index()
max_timestamp_u.columns = ['user_id', 'max_time_stamp']
max_timestamp_u.user_id=max_timestamp_u.user_id.astype('int32')

train_df['lagtime'] = train_df.groupby('user_id')['timestamp'].shift()

max_timestamp_u2 = train_df[['user_id','lagtime']].groupby(['user_id']).agg(['max']).reset_index()
max_timestamp_u2.columns = ['user_id', 'max_time_stamp2']
max_timestamp_u2.user_id=max_timestamp_u2.user_id.astype('int32')

train_df['lagtime']=train_df['timestamp']-train_df['lagtime']

train_df['lagtime']=train_df['lagtime']/(10000)

train_df.lagtime[train_df.lagtime>20000]=20000
train_df['lagtime']
train_df['lagtime2'] = train_df.groupby('user_id')['timestamp'].shift(2)
max_timestamp_u3 = train_df[['user_id','lagtime2']].groupby(['user_id']).agg(['max']).reset_index()
max_timestamp_u3.columns = ['user_id', 'max_time_stamp3']
max_timestamp_u3.user_id=max_timestamp_u3.user_id.astype('int32')

train_df['lagtime2']=train_df['timestamp']-train_df['lagtime2']
train_df['lagtime2']=train_df['lagtime2']/(10000)
#train_df.lagtime2=train_df.lagtime2.astype('float32')
train_df.lagtime2[train_df.lagtime2>20000]=20000
train_df['lagtime2'].max()
train_df['lagtime3'] = train_df.groupby('user_id')['timestamp'].shift(3)
train_df['lagtime3']=train_df['lagtime3']/(10000)
#train_df.lagtime3=train_df.lagtime3.astype('float32')
train_df.lagtime3[train_df.lagtime3>20000]=20000
max_timestamp_u_dict=max_timestamp_u.set_index('user_id').to_dict()
max_timestamp_u_dict2=max_timestamp_u2.set_index('user_id').to_dict()
max_timestamp_u_dict3=max_timestamp_u3.set_index('user_id').to_dict()

del max_timestamp_u
del max_timestamp_u2
del max_timestamp_u3
gc.collect()
prior_question_elapsed_time_mean=train_df['prior_question_elapsed_time'].mean()
train_df['prior_question_elapsed_time'].fillna(prior_question_elapsed_time_mean, inplace=True)
lagtime_mean=train_df['lagtime'].mean()
train_df['lagtime'].fillna(0, inplace=True)
train_df.lagtime=train_df.lagtime.astype('int16')
lagtime_mean2=train_df['lagtime2'].mean()
train_df['lagtime2'].fillna(0, inplace=True)
train_df.lagtime2=train_df.lagtime2.astype('int16')
lagtime_mean3=train_df['lagtime3'].mean()
train_df['lagtime3'].fillna(0, inplace=True)
train_df.lagtime3=train_df.lagtime3.astype('int16')
train_df['prior_question_elapsed_time']=train_df['prior_question_elapsed_time']/(1000)
train_df.prior_question_elapsed_time=train_df.prior_question_elapsed_time.astype('int16')
train_df.prior_question_had_explanation=train_df.prior_question_had_explanation.astype('int8')
questions_df = pd.read_csv(
    '../input/riiid-test-answer-prediction/questions.csv', 
    usecols=[0, 1,3,4],
    dtype={'question_id': 'int16','bundle_id': 'int16', 'part': 'int8','tags': 'str'}
)
tag = questions_df["tags"].str.split(" ", n = 10, expand = True)
tag.columns = ['tags1','tags2','tags3','tags4','tags5','tags6']
#
#有的tag没值，赋0值
tag.fillna(0, inplace=True)
tag = tag.astype('int16')
questions_df =  pd.concat([questions_df,tag],axis=1).drop(['tags'],axis=1)

questions_df.rename(columns={'question_id':'content_id'}, inplace=True)
questions_df['tags2'].max()
#train_df=train_df[0:3300*10000]
train_df = pd.merge(train_df, questions_df, on='content_id', how='left',right_index=True)
train_df['answer_shift'] = train_df.groupby('user_id')['answered_correctly'].shift()
train_df['answer_shift'].fillna(2, inplace=True)
train_df.answer_shift=train_df.answer_shift.astype('int8')
train_df.dtypes
#for inference\train
train_df=train_df.groupby('user_id').tail(50)
len(train_df)
##content_id have 0，so add 1
train_df['content_id'] += 1
# train_df['answered_correctly'] += 1

#train_df = train_df[train_df.content_type_id == False]
if train_flag:
    valid_df=train_df[1300*10000:]
    train_df=train_df[0:1200*10000]#less 100
else:
    valid_df=train_df[1200*10000:1300*10000]
    test_df=train_df[1200*10000:1201*10000]
train_dataset = train_df.groupby('user_id').apply(lambda r:[
            r['content_id'].values.tolist(),
            r['answered_correctly'].values.tolist(),
            r['prior_question_elapsed_time'].values.tolist(),
            r['prior_question_had_explanation'].values.tolist(),
            r['user_lecture_lv'].values.tolist(),
            r['lagtime'].values.tolist(),
            r['lagtime2'].values.tolist(),
            r['answer_shift'].values.tolist(),
            r['part'].values.tolist(),
            r['tags1'].values.tolist(),
            r['tags2'].values.tolist()])

valid_dataset = valid_df.groupby('user_id').apply(lambda r:[
            r['content_id'].values.tolist(),
            r['answered_correctly'].values.tolist(),
            r['prior_question_elapsed_time'].values.tolist(),
            r['prior_question_had_explanation'].values.tolist(),
            r['user_lecture_lv'].values.tolist(),
            r['lagtime'].values.tolist(),
            r['lagtime2'].values.tolist(),
            r['answer_shift'].values.tolist(),
            r['part'].values.tolist(),
            r['tags1'].values.tolist(),
            r['tags2'].values.tolist()])
del train_df
del valid_df
print('len(train_dataset):',len(train_dataset))
def getlimitdata(x):    
    seq_len = len(x)
    r = np.zeros(MAX_LENGTH, dtype=np.float32)

    if seq_len >= MAX_LENGTH:          
        r[:] = x[-MAX_LENGTH:]
       
    else:        
        #r[-seq_len:] = x
        r[0:seq_len] = x
            
    return r.tolist()

def getlimitdata2(x,start_token):#response
    seq_len = len(x)
    r = np.zeros(MAX_LENGTH, dtype=np.float32)

    if seq_len >= MAX_LENGTH:          
        r[:] = x[-MAX_LENGTH:]
        r[0:1]=start_token #start token
       
    else:        
        #r[-seq_len:] = x
        r[0:seq_len] = x
        r[0:1]=start_token #start token
            
    return r.tolist()
#MAX_LENGTH=50
getlimitdata2([1,3,3,4,5],8)
def getslice(dataset):
    content_ids=[]# r['content_id'].values.tolist(),
    answered_correctlys=[]# r['answered_correctly'].values.tolist(),
    prior_question_elapsed_times=[]# r['prior_question_elapsed_time'].values.tolist(),
    prior_question_had_explanations=[]# r['prior_question_had_explanation'].values.tolist(),
    user_lecture_lvs=[]# r['user_lecture_lv'].values.tolist(),
    lagtimes=[]# r['lagtime'].values.tolist(),
    lagtime2s=[]# r['lagtime2'].values.tolist(),
    answer_shifts=[]# r['answer_shift'].values.tolist(),
    parts=[]# r['part'].values.tolist(),
    tags1s=[]# r['tags1'].values.tolist(),
    tags2s=[]# r['tags2'].values.tolist()])

    i=0
    for v in dataset.values:
        #print(v[0])
        content_ids.append(getlimitdata(v[0]))
        answered_correctlys.append(getlimitdata(v[1]))
        prior_question_elapsed_times.append(getlimitdata2(v[2],elapsed_time_count-1))#response decoder
        prior_question_had_explanations.append(getlimitdata2(v[3],2))#response decoder
        user_lecture_lvs.append(getlimitdata(v[4]))
        lagtimes.append(getlimitdata2(v[5],lagtime_count-1))#response decoder
        lagtime2s.append(getlimitdata2(v[6],lagtime_count-1))#response decoder
        answer_shifts.append(getlimitdata2(v[7],2))#response decoder
        parts.append(getlimitdata(v[8]))
        tags1s.append(getlimitdata(v[9]))
        tags2s.append(getlimitdata(v[10]))
    #     if i==1:
    #         break
    #     i=i+1
    
    return content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s
%%time
if train_flag:
    train_dataset = tf.data.Dataset.from_tensor_slices((getslice(train_dataset)))

valid_dataset = tf.data.Dataset.from_tensor_slices((getslice(valid_dataset)))
# del content_ids
# del answered_correctlys
# del prior_question_elapsed_times
# del prior_question_had_explanations
# del user_lecture_lvs
# del lagtimes
# del lagtime2s
# del answer_shifts
# del parts
# del tags1s
# del tags2s
if train_flag:
    train_dataset = train_dataset.cache()
    train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(BATCH_SIZE)
    train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
valid_dataset = valid_dataset.padded_batch(2000)
content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s =next(iter(valid_dataset))
content_ids
user_lecture_lvs
user_lecture_lv_embeddings= tf.keras.layers.Dense(d_model, use_bias=False)(user_lecture_lvs)
user_lecture_lv_embeddings
content_embeddings = tf.keras.layers.Embedding(input_vocab_size, d_model)(content_ids)
"""
# **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])
    #print(angle_rads)
    pos_encoding = angle_rads[np.newaxis, ...]
    
    #pos_encoding=np.repeat(pos_encoding,BATCH_SIZE,0)#在0 这个维度扩展batchsize遍
    
    return tf.cast(pos_encoding, dtype=tf.float32)
x=positional_encoding(50,d_model)
x
x[:28,:,]
# import torch
# from torch import nn
# seq = torch.arange(50).unsqueeze(0)
# pos=nn.Embedding(50,128)(seq)
# pos
"""
# **Mask**
"""
def create_padding_mask(seq):
    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)#

    # add extra dimensions to 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):
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)
create_look_ahead_mask(5)
# import torch
# seq_len=5
# torch.triu(torch.ones(seq_len,seq_len),diagonal=1).to(dtype=torch.bool)
"""
# **Scaled dot product attention**
"""
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.math.sqrt(dk)

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

    # 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
"""
# **Multi-head attention**
"""
class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.dense = tf.keras.layers.Dense(d_model)

    def split_heads(self, x, batch_size):
       
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, v, k, q, mask):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)  # (batch_size, seq_len, d_model)
        k = self.wk(k)  # (batch_size, seq_len, d_model)
        v = self.wv(v)  # (batch_size, seq_len, d_model)

        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
        scaled_attention, attention_weights = scaled_dot_product_attention(
            q, k, v, mask)

        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

        concat_attention = tf.reshape(scaled_attention, 
                                      (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

        return output, attention_weights
"""
# **Point wise feed forward network**
"""
def point_wise_feed_forward_network(d_model, dff):
    return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
    ])
"""
# **Encoder and decoder**
"""
class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model, dff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask):

        attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)

        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

        return out2
class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
               maximum_position_encoding, rate=0.1):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding1 = tf.keras.layers.Embedding(input_vocab_size, d_model)
        self.embedding2 =tf.keras.layers.Embedding(user_lecture_count, d_model)
        self.embedding3 =tf.keras.layers.Embedding(8, d_model)
        self.embedding4 = tf.keras.layers.Embedding(188, d_model)
        self.embedding5 = tf.keras.layers.Embedding(188, d_model)
        self.pos_encoding = positional_encoding(maximum_position_encoding, 
                                                self.d_model)


        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                           for _ in range(num_layers)]

        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask):
       
        content_ids=x[0]
        content_embeddings = self.embedding1(content_ids)
        #user_lecture_lv_embeddings = tf.keras.layers.Dense(d_model, use_bias=False)(x[1])
        user_lecture_lv_embeddings = self.embedding2(x[1])
       
        part_embeddings = self.embedding3(x[2])
        tags1_embeddings = self.embedding4(x[3])
        tags2_embeddings = self.embedding5(x[4])
        
        seq_len = tf.shape(content_ids)[1]
        batch=tf.shape(content_ids)[0]
        
        pos= self.pos_encoding[:, :seq_len, :]
        
        x=content_embeddings+part_embeddings+tags1_embeddings
        #x=x+tags2_embeddings+user_lecture_lv_embeddings
       
        # Add embeddings
#         x = tf.keras.layers.Add()([            
#             content_embeddings,
#             user_lecture_lv_embeddings,
#             part_embeddings,
#             tags1_embeddings,
#             tags2_embeddings
#             #pos
#         ])
        
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x+=pos
        # 
        #x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
     

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x  # (batch_size, input_seq_len, d_model)
class DecoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(DecoderLayer, self).__init__()

        self.mha1 = MultiHeadAttention(d_model, num_heads)
        self.mha2 = MultiHeadAttention(d_model, num_heads)

        self.ffn = point_wise_feed_forward_network(d_model, dff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        self.dropout3 = tf.keras.layers.Dropout(rate)


    def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
    # enc_output.shape == (batch_size, input_seq_len, d_model)

        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layernorm1(attn1 + x)

        attn2, attn_weights_block2 = self.mha2(
            enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)

        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)

        return out3, attn_weights_block1, attn_weights_block2
class Decoder(tf.keras.layers.Layer):
    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,
               maximum_position_encoding, rate=0.1):
        super(Decoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding1 = tf.keras.layers.Embedding(3, d_model)
        self.embedding2 = tf.keras.layers.Embedding(3, d_model)
        self.embedding3 = tf.keras.layers.Embedding(elapsed_time_count, d_model)
        self.embedding4 = tf.keras.layers.Embedding(lagtime_count, d_model)
        self.embedding5 = tf.keras.layers.Embedding(lagtime_count, d_model)
        
        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)

        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) 
                           for _ in range(num_layers)]
        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
        #encode_inputs=[content_ids,user_lecture_lvs,parts,tags1s,tags2s]
        #decode_inputs=[answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,lagtimes,lagtime2s,answer_shifts]
     
        prior_question_elapsed_times=x[1]     
        prior_question_had_explanations=x[2]   
        lagtimes=x[3]       
        lagtime2s=x[4]       
        answer_shifts=x[5]
        
        
        answer_shifts_embeddings = self.embedding1(answer_shifts)
        had_explanations_embeddings = self.embedding2(prior_question_had_explanations)
#         elapsed_time_embeddings = tf.keras.layers.Dense(d_model, use_bias=False)(prior_question_elapsed_times)
#         lagtimes_embeddings = tf.keras.layers.Dense(d_model, use_bias=False)(lagtimes)
#         lagtime2s_embeddings = tf.keras.layers.Dense(d_model, use_bias=False)(lagtime2s)
#         lagtime_count=20002
#         elapsed_time_count=500
        elapsed_time_embeddings = self.embedding3(prior_question_elapsed_times)
        lagtimes_embeddings = self.embedding4(lagtimes)
        lagtime2s_embeddings = self.embedding5(lagtime2s)
        #lagtime3s_embeddings = tf.keras.layers.Dense(d_model, use_bias=False)(lagtime3s)
        
#         elapsed_time_embeddings=prior_question_elapsed_times
#         lagtimes_embeddings=lagtimes
#         lagtime2s_embeddings=lagtime2s
        
        seq_len = tf.shape(answer_shifts)[1]
        batch=tf.shape(answer_shifts)[0]
        attention_weights = {}
        
        pos = self.pos_encoding[:, :seq_len, :]
        
        x=answer_shifts_embeddings+had_explanations_embeddings
        #x=x+elapsed_time_embeddings+lagtimes_embeddings+lagtime2s_embeddings
        
#         x = tf.keras.layers.Add()([            
#             answer_shifts_embeddings,
#             had_explanations_embeddings,
#             elapsed_time_embeddings,
#             lagtimes_embeddings,
#             lagtime2s_embeddings
#             #lagtime3s_embeddings,
#             #pos
#         ])
        
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x+=pos   
        #x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                                 look_ahead_mask, padding_mask)

            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights
"""
# ** Transformer**
"""
class Transformer(tf.keras.Model):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
               target_vocab_size, pe_input, pe_target, rate=0.1):
        super(Transformer, self).__init__()

        self.encoder = Encoder(num_layers, d_model, num_heads, dff, 
                               input_vocab_size, pe_input, rate)

        self.decoder = Decoder(num_layers, d_model, num_heads, dff, 
                               target_vocab_size, pe_target, rate)

        self.final_layer = tf.keras.layers.Dense(target_vocab_size,activation='softmax')#

    def call(self, encode_inputs, decode_inputs, training, enc_padding_mask, 
           look_ahead_mask, dec_padding_mask):
       
        enc_output = self.encoder(encode_inputs, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)

        # dec_output.shape == (batch_size, tar_seq_len, d_model)
        dec_output, attention_weights = self.decoder(
            decode_inputs, enc_output, training, look_ahead_mask, dec_padding_mask)

        final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)
        

        return final_output, attention_weights
"""
# **Optimizer**
"""
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()

        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)

        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps ** -1.5)

        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
learning_rate = CustomSchedule(d_model)
#SAINT:We use the Adam optimizer  with lr =0:001;b1 = 0:9;b2 = 0:999 and epsilon = 1e􀀀8.
optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, 
                                     epsilon=1e-9)
"""
# **Loss and metrics**
"""
# from_logits=False，output为经过softmax输出的概率值。
# from_logits=True，output为经过网络直接输出的 logits张量。
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=False, reduction='none')

def loss_function(mask,real, pred):
    #mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)


def accuracy_function(mask,real, pred):
    real = tf.cast(real, dtype=tf.int64)
    accuracies = tf.equal(real, tf.argmax(pred, axis=2))

    #mask = tf.math.logical_not(tf.math.equal(real, 0))
    accuracies = tf.math.logical_and(mask, accuracies)

    accuracies = tf.cast(accuracies, dtype=tf.float32)
    mask = tf.cast(mask, dtype=tf.float32)
    return tf.reduce_sum(accuracies)/tf.reduce_sum(mask)

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.Mean(name='train_accuracy')
val_loss = tf.keras.metrics.Mean(name='val_loss')
val_accuracy = tf.keras.metrics.Mean(name='val_accuracy')
"""
# **Training and checkpointing**
"""
transformer = Transformer(num_layers, d_model, num_heads, dff,
                          input_vocab_size, target_vocab_size, 
                          pe_input=50, 
                          pe_target=50,
                          rate=dropout_rate)
# def create_masks(inp, tar):
#     # 
#     enc_padding_mask = create_look_ahead_mask(tf.shape(inp)[1]) 
#     dec_padding_mask = create_look_ahead_mask(tf.shape(inp)[1])
#     #dec_padding_mask=dec_padding_mask[1:,]
#     combined_mask = create_look_ahead_mask(tf.shape(tar)[1])

#     return enc_padding_mask, combined_mask, dec_padding_mask
def create_masks(inp, tar):
  # Encoder padding mask
    enc_padding_mask = create_padding_mask(inp)
    look_ahead_mask = create_look_ahead_mask(tf.shape(inp)[1])
    enc_padding_mask = tf.maximum(enc_padding_mask, look_ahead_mask)

    # Used in the 2nd attention block in the decoder.
    # This padding mask is used to mask the encoder outputs.
    dec_padding_mask = create_padding_mask(inp)
    look_ahead_mask = create_look_ahead_mask(tf.shape(inp)[1])
    dec_padding_mask = tf.maximum(dec_padding_mask, look_ahead_mask)
    
    # 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(inp)[1])
    dec_target_padding_mask = create_padding_mask(inp)
    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

    return enc_padding_mask, combined_mask, dec_padding_mask
checkpoint_path = "../input/riiid-saint-transformer-6-1/checkpoints/train"

ckpt = tf.train.Checkpoint(transformer=transformer,
                           optimizer=optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# 
if ckpt_manager.latest_checkpoint:
    ckpt.restore(ckpt_manager.latest_checkpoint)
    print ('Latest checkpoint restored!!')
    
    
#save path
checkpoint_path = "./checkpoints/train"

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)
# The @tf.function trace-compiles train_step into a TF graph for faster
# execution. The function specializes to the precise shape of the argument
# tensors. To avoid re-tracing due to the variable sequence lengths or variable
# batch sizes (the last batch is smaller), use input_signature to specify
# more generic shapes.

# train_step_signature = [
#     tf.TensorSpec(shape=(None, None), dtype=tf.int64),
#     tf.TensorSpec(shape=(None, None), dtype=tf.int64),
# ]

# @tf.function(input_signature=train_step_signature)
def train_step(content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s):
#     tar_inp = answered_correctlys[:, :-1]
#     tar_real = answered_correctlys[:, 1:]    
    tar_inp = answer_shifts
    tar_real = answered_correctlys
    #sigmoid,just last tar
    #tar_real=tar[: ,-1]
    encode_inputs=[content_ids,user_lecture_lvs,parts,tags1s,tags2s]
    decode_inputs=[answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,lagtimes,lagtime2s,answer_shifts]
                                      
     
    mask = tf.math.logical_not(tf.math.equal(content_ids, 0))#content_ids[:, 1:]

    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(content_ids, tar_inp)

    with tf.GradientTape() as tape:
        predictions, _ = transformer(encode_inputs, decode_inputs, 
                                     True, 
                                     enc_padding_mask, 
                                     combined_mask, 
                                     dec_padding_mask)
        loss = loss_function(mask,tar_real, predictions)

    gradients = tape.gradient(loss, transformer.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

    train_loss(loss)
    train_accuracy(accuracy_function(mask,tar_real, predictions))
    
def val_step(content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s):
#     tar_inp = answered_correctlys[:, :-1]
#     tar_real = answered_correctlys[:, 1:]    
    tar_inp = answer_shifts
    tar_real = answered_correctlys
    #sigmoid,just last tar
    #tar_real=tar[: ,-1]
    encode_inputs=[content_ids,user_lecture_lvs,parts,tags1s,tags2s]
    decode_inputs=[answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,lagtimes,lagtime2s,answer_shifts]
                                      
     
    mask = tf.math.logical_not(tf.math.equal(content_ids, 0))#content_ids[:, 1:]

    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(content_ids, tar_inp)

    
    predictions, _ = transformer(encode_inputs, decode_inputs, 
                                 False, 
                                 enc_padding_mask, 
                                 combined_mask, 
                                 dec_padding_mask)
    loss = loss_function(mask,tar_real, predictions)   

    val_loss(loss)
    val_accuracy(accuracy_function(mask,tar_real, predictions))
print('start train....')
best_val_accuracy=0
persistent=0

if train_flag:
    for epoch in range(EPOCHS):
        start = time.time()

        train_loss.reset_states()
        train_accuracy.reset_states()
        val_loss.reset_states()
        val_accuracy.reset_states()

        # inp -> portuguese, tar -> english
        #
        for (batch, (content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s)) in enumerate(train_dataset):
            train_step(content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s)

            if batch % 5 == 0:
                print ('Epoch {} Batch {} train_Loss {:.4f} train_accuracy {:.4f}'.format(
                  epoch + 1, batch, train_loss.result(), train_accuracy.result()))

        if (epoch + 1) % 5 == 0:
            ckpt_save_path = ckpt_manager.save()
            print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,
                                                                 ckpt_save_path))

        for (batch,(v_content_ids,v_answered_correctlys,v_prior_question_elapsed_times,v_prior_question_had_explanations,v_user_lecture_lvs,v_lagtimes,v_lagtime2s,v_answer_shifts,v_parts,v_tags1s,v_tags2s)) in enumerate(valid_dataset):
            val_step(v_content_ids,v_answered_correctlys,v_prior_question_elapsed_times,v_prior_question_had_explanations,v_user_lecture_lvs,v_lagtimes,v_lagtime2s,v_answer_shifts,v_parts,v_tags1s,v_tags2s)
            if batch % 5 == 0:
                print ('Epoch {} Batch {}  val_Loss {:.4f} val_accuracy {:.4f}'.format(
                  epoch + 1, batch,val_loss.result(), val_accuracy.result()))

        print ('Epoch {} train_Loss {:.4f} train_accuracy {:.4f} val_Loss {:.4f} val_accuracy {:.4f}'.format(epoch + 1, 
                                                    train_loss.result(), 
                                                    train_accuracy.result(),
                                                    val_loss.result(), 
                                                    val_accuracy.result()))
        if val_accuracy.result()>best_val_accuracy:
            best_val_accuracy=val_accuracy.result()
            persistent=0
            print('best val_accuracy:',best_val_accuracy)
        else:
            persistent+=1
            print('val_accuracy not improve:',persistent)
        if persistent>=4:
            print('earling stop...................')
            break

        print ('Time taken for 1 epoch: {} secs\n'.format(time.time() - start))
"""
# **Inference**
"""
auc_sum=0
for (batch,(content_ids,answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,user_lecture_lvs,lagtimes,lagtime2s,answer_shifts,parts,tags1s,tags2s)) in enumerate(valid_dataset):
    print(batch)
    label=answered_correctlys[: ,-1]
    label
    #input
    tar_inp = answer_shifts
    encode_inputs=[content_ids,user_lecture_lvs,parts,tags1s,tags2s]
    decode_inputs=[answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,lagtimes,lagtime2s,answer_shifts]
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
            content_ids, tar_inp)
    predictions, attention_weights = transformer(encode_inputs, 
                                                     decode_inputs,
                                                     False,
                                                     enc_padding_mask,
                                                     combined_mask,
                                                     dec_padding_mask)
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)
    predictions=predictions[:,:,-2]
    auc=roc_auc_score(label, predictions)
    auc_sum+=auc
    print('valid-auc:',auc,'auc_avg:',auc_sum/(batch+1) )
    if batch ==10:
        break
# train_dataset = train_df.groupby('user_id').apply(lambda r:[
#             r['content_id'].values.tolist(),
#             r['answered_correctly'].values.tolist(),
#             r['prior_question_elapsed_time'].values.tolist(),
#             r['prior_question_had_explanation'].values.tolist(),
#             r['user_lecture_lv'].values.tolist(),
#             r['lagtime'].values.tolist(),
#             r['lagtime2'].values.tolist(),
#             r['answer_shift'].values.tolist(),
#             r['part'].values.tolist(),
#             r['tags1'].values.tolist(),
#             r['tags2'].values.tolist()])
def getpredatas(test_df):   
    content_ids=[]
    answered_correctlys=[]    
    prior_question_elapsed_times=[]
    prior_question_had_explanations=[]
    user_lecture_lvs=[]
    lagtimes=[]
    lagtime2s=[]
    answer_shifts=[]
    parts=[]
    tags1s=[]
    tags2s=[]
    
    test_user_ids = test_df['user_id'].values
    test_content_ids = test_df['content_id'].values
    test_elapsed_times = test_df['prior_question_elapsed_time'].values
    test_had_explanations = test_df['prior_question_had_explanation'].values
    test_user_lecture_lvs = test_df['user_lecture_lv'].values
    test_lagtimes = test_df['lagtime'].values
    test_lagtime2s = test_df['lagtime2'].values
    #test_lagtime3s = test_df['lagtime3'].values
    test_parts = test_df['part'].values
    test_tags1s = test_df['tags1'].values
    test_tags2s = test_df['tags2'].values
    for user_id, content_id,prior_question_elapsed_time,prior_question_had_explanation,lecture_lv,lagtime,lagtime2,part,tags1,tags2 in zip(test_user_ids,test_content_ids,test_elapsed_times,test_had_explanations,test_user_lecture_lvs,test_lagtimes,test_lagtime2s,test_parts,test_tags1s,test_tags2s):
        if user_id in train_dataset.keys():
            test_content=train_dataset[user_id][0]
            test_content=np.concatenate((test_content,[content_id+1]))            
            
            #answered_correctly has no future.
            test_answered_correctly=train_dataset[user_id][1]
            test_answered_correctly=np.concatenate((test_answered_correctly,[2]))#no use          
            
            test_elapsed_time=train_dataset[user_id][2]
            test_elapsed_time=np.concatenate((test_elapsed_time,[prior_question_elapsed_time]))
            #
            
            test_had_explanation=train_dataset[user_id][3]
            test_had_explanation=np.concatenate((test_had_explanation,[prior_question_had_explanation]))           
            
            test_user_lecture_lv=train_dataset[user_id][4]
            test_user_lecture_lv=np.concatenate((test_user_lecture_lv,[lecture_lv]))           
            
            test_lagtime=train_dataset[user_id][5]         
            test_lagtime=np.concatenate((test_lagtime,[lagtime]))          
            
            test_lagtime2=train_dataset[user_id][6]         
            test_lagtime2=np.concatenate((test_lagtime2,[lagtime2]))           
            
            test_answer_shifts=train_dataset[user_id][7]# when prior_test_df, update                    
            
            #,part,tags1,tags2
            test_part=train_dataset[user_id][8]         
            test_part=np.concatenate((test_part,[part]))            
            
            test_tags1=train_dataset[user_id][9]         
            test_tags1=np.concatenate((test_tags1,[tags1]))           
            
            test_tags2=train_dataset[user_id][10]         
            test_tags2=np.concatenate((test_tags2,[tags2]))            
        else:
            #test_content = np.zeros(MAX_LENGTH-1, dtype=int)
            test_content=np.array([content_id+1])
            
            #test_answered_correctly = np.zeros(MAX_LENGTH-1, dtype=int)
            test_answered_correctly=np.array([2])#no use
            
            #test_elapsed_time = np.zeros(MAX_LENGTH-1, dtype=int)
            test_elapsed_time=np.array([prior_question_elapsed_time])
            
            #test_had_explanation = np.zeros(MAX_LENGTH-1, dtype=int)
            test_had_explanation=np.array([prior_question_had_explanation])            
            
            #test_user_lecture_lv = np.zeros(MAX_LENGTH-1, dtype=int)
            test_user_lecture_lv=np.array([lecture_lv])
            
            #test_lagtime = np.zeros(MAX_LENGTH-1, dtype=int)
            test_lagtime=np.array([lagtime])
            
            #test_lagtime2 = np.zeros(MAX_LENGTH-1, dtype=int)
            test_lagtime2=np.array([lagtime2])
            
            
            test_answer_shifts=np.array([2])
            
            #test_part = np.zeros(MAX_LENGTH-1, dtype=int)
            test_part=np.array([part])
            
            #test_tags1 = np.zeros(MAX_LENGTH-1, dtype=int)
            test_tags1=np.array([tags1])
            
            #test_tags2 = np.zeros(MAX_LENGTH-1, dtype=int)
            test_tags2=np.array([tags2])
        
        train_dataset[user_id]=[test_content,test_answered_correctly,test_elapsed_time,test_had_explanation,test_user_lecture_lv,test_lagtime,test_lagtime2,test_answer_shifts,test_part,test_tags1,test_tags2]
        
        test_content=getlimitdata(test_content)        
        content_ids.append(test_content)      
        
        #answered_correctly has no future.so max-1  dimension
        test_answered_correctly=getlimitdata(test_answered_correctly)
        answered_correctlys.append(test_answered_correctly)
        
        test_elapsed_time=getlimitdata2(test_elapsed_time,elapsed_time_count-1)        
        prior_question_elapsed_times.append(test_elapsed_time)
        
        test_had_explanation=getlimitdata2(test_had_explanation,2)
        prior_question_had_explanations.append(test_had_explanation)
        
        test_user_lecture_lv=getlimitdata(test_user_lecture_lv)
        user_lecture_lvs.append(test_user_lecture_lv)
        
        test_lagtime=getlimitdata2(test_lagtime,lagtime_count-1)
        lagtimes.append(test_lagtime)
            
        test_lagtime2=getlimitdata2(test_lagtime2,lagtime_count-1)
        lagtime2s.append(test_lagtime2)
            
        test_answer_shifts=getlimitdata2(test_answer_shifts,2)
        answer_shifts.append(test_answer_shifts)
            
        test_part=getlimitdata(test_part)
        parts.append(test_part)
            
        test_tags1=getlimitdata(test_tags1)
        tags1s.append(test_tags1)
            
        test_tags2=getlimitdata(test_tags2)
        tags2s.append(test_tags2)
            
        
    
    content_ids=tf.convert_to_tensor(content_ids)
    user_lecture_lvs=tf.convert_to_tensor(user_lecture_lvs)
    parts=tf.convert_to_tensor(parts)
    tags1s=tf.convert_to_tensor(tags1s)
    tags2s=tf.convert_to_tensor(tags2s)
    answered_correctlys=tf.convert_to_tensor(answered_correctlys)
    prior_question_elapsed_times=tf.convert_to_tensor(prior_question_elapsed_times)
    prior_question_had_explanations=tf.convert_to_tensor(prior_question_had_explanations)
    lagtimes=tf.convert_to_tensor(lagtimes)
    lagtime2s=tf.convert_to_tensor(lagtime2s)
    answer_shifts=tf.convert_to_tensor(answer_shifts)

    encode_inputs=[content_ids,user_lecture_lvs,parts,tags1s,tags2s]
    decode_inputs=[answered_correctlys,prior_question_elapsed_times,prior_question_had_explanations,lagtimes,lagtime2s,answer_shifts]
    return encode_inputs,decode_inputs
encode_inputs,decode_inputs =getpredatas(test_df[0:10])
content_ids=encode_inputs[0]
tar_inp=decode_inputs[5]
enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
        content_ids, tar_inp)
predictions, attention_weights = transformer(encode_inputs, 
                                                 decode_inputs,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)
predictions=predictions[:,:,-2]
np.squeeze(predictions.numpy())
import riiideducation
env = riiideducation.make_env()
iter_test = env.iter_test()
prior_test_df = None
%%time
target = 'answered_correctly'
for (test_df, sample_prediction_df) in iter_test: 
    if prior_test_df is not None:
        prior_test_df[target] = eval(test_df['prior_group_answers_correct'].iloc[0])
        prior_test_df = prior_test_df[prior_test_df[target] != -1].reset_index(drop=True)
        user_ids = prior_test_df['user_id'].values
        content_ids = prior_test_df['content_id'].values
        targets = prior_test_df[target].values 
        for user_id, answered_correctly,content_id in zip(user_ids,targets,content_ids):
            #when train ,content_id add 1
            if user_id in train_dataset.keys():
                #train_dataset[user_id][0]=np.concatenate((train_dataset[user_id][0],[content_id+1]))
                train_dataset[user_id][1]=np.concatenate((train_dataset[user_id][1],[answered_correctly]))
#             else:
# #                 train_dataset[user_id][0]=[content_id+1]
# #                 train_dataset[user_id][1]=[answered_correctly]
#                 train_df[user_id]=[np.array([content_id+1]),np.array([answered_correctly])]
             
            
        
    prior_test_df = test_df.copy()
    
    question_len=len( test_df[test_df['content_type_id'] == 0])
    
    user_lecture_sum = np.zeros(question_len, dtype=np.int16)
    user_lecture_count = np.zeros(question_len, dtype=np.int16)
    lagtime = np.zeros(question_len, dtype=np.float32)
    lagtime2 = np.zeros(question_len, dtype=np.float32)
    lagtime3 = np.zeros(question_len, dtype=np.float32)
    
    i=0
    for j, (user_id,content_type_id,timestamp) in enumerate(zip(test_df['user_id'].values,test_df['content_type_id'].values,test_df['timestamp'].values)):
        user_lecture_sum_dict[user_id] += content_type_id
        user_lecture_count_dict[user_id] += 1
        if(content_type_id==0):#question    
            user_lecture_sum[i] = user_lecture_sum_dict[user_id]
            user_lecture_count[i] = user_lecture_count_dict[user_id]
            
            if user_id in max_timestamp_u_dict['max_time_stamp'].keys():
                lagtime[i]=timestamp-max_timestamp_u_dict['max_time_stamp'][user_id]
                if(max_timestamp_u_dict2['max_time_stamp2'][user_id]==lagtime_mean2):#第二次也要赋值平均值
                    lagtime2[i]=lagtime_mean2
                    lagtime3[i]=lagtime_mean3
                    #max_timestamp_u_dict3['max_time_stamp3'].update({user_id:lagtime_mean3})
                else:
                    lagtime2[i]=timestamp-max_timestamp_u_dict2['max_time_stamp2'][user_id]
                    if(max_timestamp_u_dict3['max_time_stamp3'][user_id]==lagtime_mean3):
                        lagtime3[i]=lagtime_mean3 #lagtime_mean3第3次也要赋值平均值
                    else:
                        lagtime3[i]=timestamp-max_timestamp_u_dict3['max_time_stamp3'][user_id]
                    
                    max_timestamp_u_dict3['max_time_stamp3'][user_id]=max_timestamp_u_dict2['max_time_stamp2'][user_id]
                        
                max_timestamp_u_dict2['max_time_stamp2'][user_id]=max_timestamp_u_dict['max_time_stamp'][user_id]
                max_timestamp_u_dict['max_time_stamp'][user_id]=timestamp
#                 lagtime_means[i]=(lagtime_mean_dict[user_id]+lagtime[i])/2
#                 lagtime_mean_dict[user_id]=lagtime_means[i]
            else:
                lagtime[i]=lagtime_mean
                max_timestamp_u_dict['max_time_stamp'].update({user_id:timestamp})
                lagtime2[i]=lagtime_mean2#第一次赋值平均值
                max_timestamp_u_dict2['max_time_stamp2'].update({user_id:lagtime_mean2})
                lagtime3[i]=lagtime_mean3#第一次赋值平均值
                max_timestamp_u_dict3['max_time_stamp3'].update({user_id:lagtime_mean3})
            i=i+1 
    
    
    test_df = test_df[test_df['content_type_id'] == 0].reset_index(drop=True)
    test_df['prior_question_had_explanation'].fillna(False, inplace=True)
    test_df.prior_question_had_explanation=test_df.prior_question_had_explanation.astype('int8')
    
    test_df['user_lecture_lv'] = user_lecture_sum
    test_df["lagtime"]=lagtime
    test_df["lagtime2"]=lagtime2
    test_df["lagtime3"]=lagtime3
    
    test_df['lagtime']=test_df['lagtime']/(10000)
    test_df.lagtime[test_df.lagtime>20000]=20000
    test_df['lagtime'].fillna(0, inplace=True)
    test_df.lagtime=test_df.lagtime.astype('int16')

    test_df['lagtime2']=test_df['lagtime2']/(10000)
    test_df.lagtime2[test_df.lagtime2>20000]=20000
    test_df['lagtime2'].fillna(0, inplace=True)
    test_df.lagtime2=test_df.lagtime2.astype('int16')

    test_df['prior_question_elapsed_time'].fillna(prior_question_elapsed_time_mean, inplace=True)
    test_df['prior_question_elapsed_time']=test_df['prior_question_elapsed_time']/(1000)
    test_df.prior_question_elapsed_time=test_df.prior_question_elapsed_time.astype('int16')
    
#     user_ids = test_df['user_id'].values
#     content_ids = test_df['content_id'].values
    
    test_df=test_df.merge(questions_df.loc[questions_df.index.isin(test_df['content_id'])],
                  how='left', on='content_id', right_index=True)
 
    encode_inputs,decode_inputs=getpredatas(test_df)
    content_ids=encode_inputs[0]
    tar_inp=decode_inputs[5]

    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
        content_ids, tar_inp)
    predictions, attention_weights = transformer(encode_inputs, 
                                                 decode_inputs,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)
    predictions=predictions[:,:,-2]
    
    test_df[target]=np.squeeze(predictions.numpy())
       
    #    
    env.predict(test_df[['row_id', target]])
    