class GRU_plain(nn.Module): 
    def init(self, input_size, embedding_size, hidden_size, num_layers, has_input=True, has_output=False, output_size=None): 
        super(GRU_plain, self).init() 
        self.num_layers = num_layers 
        self.hidden_size = hidden_size 
        self.has_input = has_input 
        self.has_output = has_output
        if has_input:
            self.input = nn.Linear(input_size, embedding_size) 
            self.rnn = nn.GRU(input_size=embedding_size, hidden_size=hidden_size, num_layers=num_layers,
                              batch_first=True) 
        else:
            self.rnn = nn.GRU(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
        if has_output:
            self.output = nn.Sequential(
                nn.Linear(hidden_size, embedding_size),
                nn.ReLU(),
                nn.Linear(embedding_size, output_size)
            )
    
        self.relu = nn.ReLU()
        # initialize
        self.hidden = None  # need initialize before forward run
    
        for name, param in self.rnn.named_parameters():
            if 'bias' in name:
                nn.init.constant(param, 0.25)
            elif 'weight' in name:
                nn.init.xavier_uniform(param,gain=nn.init.calculate_gain('sigmoid'))
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.weight.data = init.xavier_uniform(m.weight.data, gain=nn.init.calculate_gain('relu'))
    
    def init_hidden(self, batch_size): 
        return Variable(torch.zeros(self.num_layers, batch_size, self.hidden_size)).cuda()
    
    def forward(self, input_raw, pack=False, input_len=None):
        if self.has_input:
            input = self.input(input_raw) 
            input = self.relu(input)
        else:
            input = input_raw
        if pack:
            input = pack_padded_sequence(input, input_len, batch_first=True) 
        output_raw, self.hidden = self.rnn(input, self.hidden)
        if pack:
            output_raw = pad_packed_sequence(output_raw, batch_first=True)[0]
        if self.has_output:
        output_raw = self.output(output_raw)
    # return hidden state at each time step
    return output_raw

# load a list of graphs
def load_graph_list(fname,is_real=True):

    for i in range(len(graph_list)):
        edges_with_selfloops = graph_list[i].selfloop_edges()
        if len(edges_with_selfloops)>0:
            graph_list[i].remove_edges_from(edges_with_selfloops)
        if is_real:
            graph_list[i] = max(nx.connected_component_subgraphs(graph_list[i]), key=len)
            graph_list[i] = nx.convert_node_labels_to_integers(graph_list[i])
        else:
            graph_list[i] = pick_connected_component_new(graph_list[i])
    return graph_list

if 'nobfs' in args.note:
    print('nobfs')
    dataset = Graph_sequence_sampler_pytorch_nobfs(graphs_train, max_num_node=args.max_num_node)
    args.max_prev_node = args.max_num_node-1
if 'barabasi_noise' in args.graph_type:
    print('barabasi_noise')
    dataset = Graph_sequence_sampler_pytorch_canonical(graphs_train,max_prev_node=args.max_prev_node)
    args.max_prev_node = args.max_num_node - 1
else:
    dataset = Graph_sequence_sampler_pytorch(graphs_train,max_prev_node=args.max_prev_node,max_num_node=args.max_num_node)
sample_strategy = torch.utils.data.sampler.WeightedRandomSampler([1.0 / len(dataset) for i in range(len(dataset))],
                                                                 num_samples=args.batch_size*args.batch_ratio, replacement=True)

elif 'GraphRNN_RNN' in args.note:
    rnn = GRU_plain(input_size=args.max_prev_node, embedding_size=args.embedding_size_rnn,
                    hidden_size=args.hidden_size_rnn, num_layers=args.num_layers, has_input=True,
                    has_output=True, output_size=args.hidden_size_rnn_output)
    output = GRU_plain(input_size=1, embedding_size=args.embedding_size_rnn_output,
                       hidden_size=args.hidden_size_rnn_output, num_layers=args.num_layers, has_input=True,
                       has_output=True, output_size=1)