import numpy as np
import pandas as pd

import tensorflow as tf
import random
from tensorflow.keras import regularizers

from scipy.stats import spearmanr
from sklearn.metrics import median_absolute_error

import sys

use_model = 'fwgcn'
options = ['mf', 'gcn', 'fwgcn']
if len(sys.argv) == 2:
    if not sys.argv[1] in options:
        print('unknown model')
        print('usage: fwgcn [mf|gcn|fwgcn]')
        quit()
    use_model = sys.argv[1]
elif len(sys.argv) > 2:
    print('unknown option')
    print('usage: fwgcn [mf|gcn|fwgcn]')
    quit()

print('running experiment using', use_model)

### parameters ###
emb_size = 20

### preparing data ###
data = pd.read_csv('edges-ncov-31.csv')
max_time = max(data['time'])
min_time = min(data['time'])
data['source'] = data['source'] - 1
data['target'] = data['target'] - 1
print('TIME from {} to {}'.format(min_time, max_time))
n_nodes = max(data['source'].to_list() + data['target'].to_list())+1
print('number of word nodes:', n_nodes)

negative_ratio = 10
training_data_by_time = {}
for i in range(min_time, max_time+1):
    data_t = data[data['time'] == i].to_numpy()
    data_t_labeled = np.append(data_t, np.ones(len(data_t)).reshape(len(data_t), 1), 1)
    data_t_neg = []
    data_d = set(data[data['time'] == i].itertuples(index=False, name=None))
    for _ in range(len(data_t)):
        neg_source = random.sample(range(n_nodes), negative_ratio-1)
        neg_target = random.sample(range(n_nodes), negative_ratio-1)
        for j in range(len(neg_source)):
            if (neg_source[j], neg_target[j], 0) in data_d or (neg_target[j], neg_source[j], 1) in data_d:
                continue
            data_t_neg.append([neg_source[j], neg_target[j], 1, j])
    data_t_neg_labeled = np.append(data_t_neg, np.zeros(len(data_t_neg)).reshape(len(data_t_neg), 1), 1)
    training_data_by_time[i] = np.concatenate([data_t_labeled, data_t_neg_labeled], axis=0)

A_by_time = {}
for i in range(min_time, max_time+1):
    A = np.zeros((n_nodes, n_nodes))
    data_t = data[data['time'] == i].to_numpy()
    for [head, tail, weight, t] in data_t:
        A[head, tail] = 1
    A_by_time[i] = A

tfidf_hour = pd.read_csv('ncov_tfidf_vs_hour.csv', header=None)
tfidf = np.zeros((n_nodes, max_time-min_time+1))
for i in range(tfidf_hour.shape[0] // 24):
    # perform normalization
    vs = tfidf_hour.iloc[i*24:(i+1)*24,:].to_numpy().mean(axis=0).tolist()
    max_v = max(vs)
    if max_v == 0:
        tfidf[:,i] = vs
        continue
    tfidf[:,i] = [v/max_v for v in vs]
    
node_embeddings = tf.keras.layers.Embedding(n_nodes, emb_size)
A_t = A_by_time[min_time]

### models ###
class MF(tf.keras.models.Model):
    def __init__(self):
        input_full = tf.keras.layers.Input(shape=(2,), dtype=tf.int32)
        head = input_full[:, 0]
        tail = input_full[:, 1]
        head_emb = node_embeddings(head)
        tail_emb = node_embeddings(tail)
        mul_vec = tf.keras.layers.multiply([head_emb, tail_emb])
        prediction = tf.keras.layers.Dense(1, activation='sigmoid')(mul_vec)
        super(MF, self).__init__(inputs=input_full, outputs=prediction)


class GCN(tf.keras.models.Model):
    def __init__(self):
        w = tf.keras.layers.Embedding(n_nodes, emb_size)

        input_full = tf.keras.layers.Input(shape=(2,), dtype=tf.int32)
        head = input_full[:, 0]
        tail = input_full[:, 1]

        head_A = tf.nn.embedding_lookup(A_t, head)
        head_A = tf.transpose(tf.cast(head_A, dtype=tf.float32))
        head_emb_l = node_embeddings(head)
        head_emb = tf.nn.embedding_lookup(head_emb_l, head)
        head_emb = tf.nn.embedding_lookup(tf.matmul(head_A, head_emb), head)
        head_emb = tf.keras.layers.multiply([head_emb, w(head)])

        tail_A = tf.nn.embedding_lookup(A_t, tail)
        tail_A = tf.transpose(tf.cast(tail_A, dtype=tf.float32))
        tail_emb_l = node_embeddings(tail)
        tail_emb = tf.nn.embedding_lookup(tail_emb_l, tail)
        tail_emb = tf.nn.embedding_lookup(tf.matmul(tail_A, tail_emb), tail)
        tail_emb = tf.keras.layers.multiply([tail_emb, w(tail)])

        mul_vec = tf.keras.layers.multiply([head_emb, tail_emb])
        prediction = tf.keras.layers.Dense(1, activation='sigmoid')(mul_vec)
        super(GCN, self).__init__(inputs=input_full, outputs=prediction)


class FWGCN(tf.keras.models.Model):
    def __init__(self):
        w = tf.keras.layers.Embedding(n_nodes, emb_size)

        input_full = tf.keras.layers.Input(shape=(2,), dtype=tf.int32)
        head = input_full[:, 0]
        tail = input_full[:, 1]

        head_A = tf.nn.embedding_lookup(A_t, head)
        head_A = tf.transpose(tf.cast(head_A, dtype=tf.float32))
        head_emb_l = node_embeddings(head)
        head_emb = tf.nn.embedding_lookup(head_emb_l, head)
        head_emb = tf.nn.embedding_lookup(tf.matmul(head_A, head_emb), head)
        head_emb = tf.keras.layers.multiply([head_emb, w(head)])

        tail_A = tf.nn.embedding_lookup(A_t, tail)
        tail_A = tf.transpose(tf.cast(tail_A, dtype=tf.float32))
        tail_emb_l = node_embeddings(tail)
        tail_emb = tf.nn.embedding_lookup(tail_emb_l, tail)
        tail_emb = tf.nn.embedding_lookup(tf.matmul(tail_A, tail_emb), tail)
        tail_emb = tf.keras.layers.multiply([tail_emb, w(tail)])

        head_tfidf = tf.nn.embedding_lookup(tfidf[:, t - min_time], head)
        tail_tfidf = tf.nn.embedding_lookup(tfidf[:, t - min_time], tail)

        head_fw = tf.keras.layers.multiply([head_emb, head_tfidf])
        tail_fw = tf.keras.layers.multiply([tail_emb, tail_tfidf])

        mul_vec = tf.keras.layers.multiply([head_emb, tail_emb])
        mul_vec_fw = tf.keras.layers.multiply([head_emb, head_tfidf, tail_emb, tail_tfidf])
        prediction = tf.keras.layers.Dense(1, activation='sigmoid')(mul_vec)
        prediction_fw = tf.keras.layers.Dense(1, activation='sigmoid')(mul_vec_fw)
        super(FWGCN, self).__init__(inputs=input_full, outputs=(prediction, prediction_fw))

### training routines ###
def trainMF(n_epoch):
    embs_at_time = {}

    model = MF()
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)

    batch_size = 1024
    for t in range(min_time, max_time+1):
        sum_loss = 0
        for epoch in range(n_epoch):
            x_train = training_data_by_time[t][:,0:2]
            y_train = training_data_by_time[t][:,4]
            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
            train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)
            for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
                with tf.GradientTape() as tape:
                    logits = model(x_batch_train, training=True)  # Logits for this minibatch
                    loss_value = loss_fn(y_batch_train[:, tf.newaxis], logits)
                    sum_loss += loss_value
                grads = tape.gradient(loss_value, model.trainable_weights)
                optimizer.apply_gradients(zip(grads, model.trainable_weights))
        embs_at_time[t] = node_embeddings.weights[0].numpy().copy()
        print('time {}, loss {}'.format(t, sum_loss/n_epoch))
    return embs_at_time

def trainGCN(n_epoch):
    embs_at_time = {}

    model = GCN()
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)

    batch_size = 1024
    for t in range(min_time, max_time+1):
        A_t = A_by_time[t]
        sum_loss = 0
        for epoch in range(n_epoch):
            x_train = training_data_by_time[t][:,0:2]
            y_train = training_data_by_time[t][:,4]
            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
            train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)
            for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
                with tf.GradientTape() as tape:
                    logits = model(x_batch_train, training=True)  # Logits for this minibatch
                    loss_value = loss_fn(y_batch_train[:, tf.newaxis], logits)
                    sum_loss += loss_value
                grads = tape.gradient(loss_value, model.trainable_weights)
                optimizer.apply_gradients(zip(grads, model.trainable_weights))
        embs_at_time[t] = node_embeddings.weights[0].numpy().copy()
        print('time {}, loss {}'.format(t, sum_loss/n_epoch))
    return embs_at_time
    
def trainFWGCN(n_epoch):
    embs_at_time = {}
    model = FWGCN()
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)
    loss_kl = tf.keras.losses.KLDivergence()

    batch_size = 1024
    for t in range(min_time, max_time+1):
        A_t = A_by_time[t]
        sum_loss, sum_loss_div = 0, 0
        for epoch in range(n_epoch):
            x_train = training_data_by_time[t][:,0:2]
            y_train = training_data_by_time[t][:,4]
            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
            train_dataset = train_dataset.shuffle(buffer_size=1024).batch(batch_size)
            for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
                with tf.GradientTape() as tape:
                    prediction, prediction_fw = model(x_batch_train, training=True)  # Logits for this minibatch
                    loss_comp = loss_fn(y_batch_train[:, tf.newaxis], prediction_fw)

                    prediction_dist = tf.concat([prediction, 1.0-prediction], 1)
                    prediction_dist_fw = tf.concat([prediction_fw, 1.0-prediction_fw], 1)
                    loss_div = loss_kl(prediction_dist_fw, prediction_dist)
                    loss_value = loss_comp + loss_div
                    sum_loss += loss_value
                    sum_loss_div += loss_div
                grads = tape.gradient(loss_value, model.trainable_weights)
                optimizer.apply_gradients(zip(grads, model.trainable_weights))
        embs_at_time[t] = node_embeddings.get_weights()[0] * tfidf[:,t-min_time].reshape((n_nodes, 1))
        print('time {}, loss {}, div loss {}'.format(t, sum_loss/n_epoch, sum_loss_div/n_epoch))
    return embs_at_time

### running representation learning ####
embs_at_time = None
if use_model == 'fwgcn':
    embs_at_time = trainFWGCN(n_epoch=5)
elif use_model == 'gcn':
    embs_at_time = trainGCN(n_epoch=5)
elif use_model == 'mf':
    embs_at_time = trainMF(n_epoch=5)
else:
    print('unknown model: ', use_model)
    quit()

### downstream task ###
print('using learned embeddings to perform downstream task: tweet popularity prediction')
embs_at_time_fix = {}
for i in range(len(embs_at_time)):
    embs_at_time_fix[i] = embs_at_time[i+1]

sns_day_to_mat = embs_at_time_fix

sns_day = {}
for i in range(len(sns_day_to_mat)):
#     vs = np.transpose(sns_day_to_mat[i], (1,0)).flatten().tolist()
    vs = sns_day_to_mat[i].flatten().tolist()
    
    max_v = max(vs)
    if max_v == 0:
        sns_day[i] = vs
        continue
    sns_day[i] = [v/max_v for v in vs]
n_words = n_nodes
word_dim = emb_size
dim_flat = n_words * word_dim
sns_vecs = np.zeros((len(sns_day_to_mat), dim_flat))
for i,t in enumerate(sns_day.keys()):
    sns_vecs[i,:] = sns_day[i]
    
train_data = np.load('popPred_ncov_bow_train.npy', allow_pickle=True)
test_data = np.load('popPred_ncov_bow_test.npy', allow_pickle=True)

tweet_vec_len = n_words * 2
text_vec_len = n_words
user_vec_len = n_words
train_x = train_data[:,2:(tweet_vec_len +5)]
train_y = train_data[:,1]
test_x = test_data[:,2:(tweet_vec_len +5)]
test_y = test_data[:,1]

input_dim = train_x.shape[1]
class MLP(tf.keras.models.Model):
    def __init__(self, layers_dim):
        input_full = tf.keras.layers.Input(shape=(input_dim,), dtype=tf.float32)
        tweet_data = input_full[:,:text_vec_len]
        user_data = input_full[:,text_vec_len:tweet_vec_len]
        tweet_info = input_full[:,tweet_vec_len:(tweet_vec_len+2)]
        day = input_full[:,(tweet_vec_len+2)]
        
        sns_data = tf.nn.embedding_lookup(sns_vecs, tf.cast(day, dtype=tf.int32))
        sns_mat = tf.keras.layers.Reshape((n_words, word_dim), input_shape=(dim_flat,))(sns_data)
        sns_mat = tf.transpose(sns_mat, [0,2,1])
        query = tweet_data
        snse_context = tf.keras.layers.Attention()([query, sns_mat])
        snse_vec = tf.reduce_mean(snse_context, 1)
        sns_vec = tf.keras.layers.multiply([snse_vec, tweet_data])
        
        mlp_vector = sns_vec
        num_layer = len(layers_dim)
        for layer in range(num_layer):
            model_layer = tf.keras.layers.Dense(
                layers_dim[layer],
                kernel_regularizer=tf.keras.regularizers.l2(1e-05),
                activation="relu")
            mlp_vector = model_layer(mlp_vector)
        
        predict_vector = tf.keras.layers.concatenate([mlp_vector, tweet_info])
        full1 = tf.keras.layers.Dense(100, activation='relu', activity_regularizer=regularizers.l2(1e-5))(predict_vector)
        dropout = tf.keras.layers.Dropout(.5, input_shape=(2,))(full1)
        prediction = tf.keras.layers.Dense(1, activation="linear")(dropout)
        super(MLP, self).__init__(inputs=input_full, outputs=prediction)

ds_model = MLP([120, 80, 50])
ds_model.compile( optimizer=tf.optimizers.Adam(learning_rate=0.001), loss='mean_squared_error')
ds_model.fit(train_x, train_y, epochs=15, batch_size=64, verbose=1, validation_split = 0.1)
predict = ds_model.predict(test_x)
mse = np.mean(abs(predict - test_y))
rmse = np.sqrt(np.mean(np.square(predict - test_y)))
spr, pval = spearmanr(predict, test_y)
medae = median_absolute_error(predict, test_y)
print(mse, rmse, medae, spr)
