# -*- coding: utf-8 -*-
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn  import GCNConv
from tqdm  import tqdm

class GlobalMaxPool1d(nn.Module):
    def __init__( self ):
        super (GlobalMaxPool1d, self ).__init__()

    def forward( self , x):
        return torch.max_pool1d(x, kernel_size = x.shape[ 2 ])


class ConvGCN(nn.Module):
    def __init__( self , dim_in, dim_hid, dim_out, dim_stance = 4 , dropout = 0.5 ):
        """
        Conversational-GCN

        :Params:
            T (int):              the number of splited timesteps for propagation graph
            dim_in (int):         结点的初始输入特征维度 k
            dim_hid (int):        默认（固定）的结点嵌入的维度（等于结点池化后的图的嵌入维度）
            dim_out (int):        模型最终的输出维度，用于分类
            num_layers (int):     LGAT的层数（邻居聚集的迭代次数）
            dropout (float):
        """
        super (ConvGCN,  self ).__init__()
        self .dropout  = dropout
        # dim_hid -> embed_size
        self .word_embeddings  = nn.Parameter(nn.init.xavier_uniform_(
            torch.zeros(dim_in, dim_hid, dtype = torch. float , device = device), gain = np.sqrt( 2.0 )), requires_grad = True )

        # BiGRU for obtaining post embedding (batch_size = 1)
        bigru_num_layers  = 1
        bigru_num_directions  = 2
        bigru_hidden  = dim_hid  / / 2
        self .BiGRU  = nn.GRU(dim_hid, bigru_hidden, bigru_num_layers, bidirectional = True )
        self .H0  = torch.zeros(bigru_num_layers  * bigru_num_directions,  1 , bigru_hidden, device = device)

        # Graph Convolution
        self .conv1  = GCNConv(dim_hid, dim_hid)
        self .conv2  = GCNConv(dim_hid, dim_stance)

        # GRU for modeling the temporal dynamics
        rnn_num_layers  = 1 # rnn_hidden = dim_hid
        self .GRU  = nn.GRU(dim_hid + dim_stance, dim_hid + dim_stance, rnn_num_layers)
        self .H1  = torch.zeros(rnn_num_layers,  1 , dim_hid + dim_stance, device = device)
       
        self .MaxPooling  = GlobalMaxPool1d()
        self .prediction_layer  = nn.Linear(dim_hid + dim_stance, dim_out)
        nn.init.xavier_normal_( self .prediction_layer.weight)

    def forward( self , words_indices, edge_indices):
        features  = []
        edge_indice  = edge_indices[ 0 ].to(device)   # only one snapshot
        words_indices  = words_indices.to(device)   # get post embedding
        for i  in range (words_indices.shape[ 0 ]):
            word_indice  = torch.nonzero(words_indices[i], as_tuple = True )[ 0 ]
            # assert word_indice.shape[0] > 0, "words must large or equal to one"
            if word_indice.shape[ 0 ]  = = 0 :
                word_indice  = torch.tensor([ 0 ], dtype = torch. long ).to(device)

            words  = self .word_embeddings.index_select( 0 , word_indice.to(device))
            _, hn  = self .BiGRU(words.unsqueeze( 1 ),  self .H0)   # (num_layers * num_directions, batch, hidden_size)
            post_embedding  = hn.flatten().unsqueeze( 0 )
            features.append(post_embedding)

        x0  = torch.cat(features, dim = 0 )
        content  = torch.clone(x0)
        x1  = self .conv1(x0, edge_indice)
        x1  = F.relu(x1)
        x1  = F.dropout(x1,  self .dropout)
        x2  = self .conv2(x1, edge_indice)
        x2  = F.relu(x2)

        x  = torch.cat((content, x2),  1 )
        gru_output, _  = self .GRU(x.unsqueeze( 1 ),  self .H1)
        z  = self .MaxPooling(gru_output.transpose( 0 ,  2 ))
        return self .prediction_layer(z.squeeze( 1 ).transpose_( 0 ,  1 ))   # 使用BCE(不需事先计算softmax)

def load_rawdata(file_path):
    """ json file, like a list of dict """
    with  open (file_path, encoding = "utf-8" ) as f:
        data  = json.loads(f.read())
    return data


def get_edges(data):
    """依据src_id's data，加载propagation network edges with relative index in a graph"""
    tweet_num  = len (data)
    mids  = [tweet[ "mid" ]  for tweet  in data]            # 使用"mid"才能找到所有转发关系
    mids_id  = {mids[i]: i  for i  in range (tweet_num)}   # mid: id  from mid to index，按顺序生成mid的id
    mid_edges  = [(mid, data[mids_id[mid]][ "parent" ])  for mid  in mids  if data[mids_id[mid]][ "parent" ] ! = None ]
    return [(mids_id[edge[ 0 ]], mids_id[edge[ 1 ]])  for edge  in mid_edges]


def get_static_edgeindex(edges):
    """ 获取传播图的edge_indices列表,要保证得到的绝对边索引对应结点特征 (for second)"""
    graph  = nx.Graph()     # len(split_edges)= T，和新增的边相关的结点就是Interacting结点
    graph.add_edges_from(edges)
    edge_index  = list (nx.adjacency_matrix(graph).nonzero())    # 结点和边的Id以及安排了，不需再处理
    return edge_index


def load_data(ids, T, thres_num, is_dynamic = False ):
    """
    依据weibo的id，加载所有的结点特征和传播结构

    :params:
        weibo_id (string):    微博id
    :returns:
        graph_list: a list of Diffusion graph objects which contain elements as follows:
            1.features:       numpy ndarray          结点数N by 特征维度k的结点初始嵌入矩阵
            2.edge_indices：  list(numpy ndarray)    对话结构的邻接矩阵
    """
    graph_list  = []
    for weibo_id  in ids:
        txt_data  = load_sptext(weibo_id)   # 从磁盘上加载以稀疏矩阵存储的text matrix
        data  = load_rawdata(data_path  + "Weibo/{0}.json" . format (weibo_id))
        edges  = get_edges(data)
        edge_indices  = get_static_edgeindex(edges)
        graph_list.append(DiffGraph(txt_data, edge_indices))

    return graph_list   # 直接把graph list中的DiffGraph拿来Pickle

class DiffGraph():
    def __init__( self , text_data, edge_indices):
        self .text_data  = text_data
        self .edge_indices  = [torch.tensor(edges, dtype = torch. long )  for edges  in edge_indices]

    def get_wordindices( self ):
        return torch.from_numpy( self .text_data.toarray())