import numpy as np
import scipy.sparse as sp
import torch
import torch.nn as nn
import networkx as nx
import time

from embed_methods.dgi.models import DGI, LogReg
from embed_methods.dgi.utils import process

def dgi(G, features):

  batch_size = 1
  nb_epochs = 10000
  patience = 20
  lr = 0.001
  l2_coef = 0.0
  drop_prob = 0.0
  hid_units = 512
  sparse = True
  nonlinearity = 'prelu' # special name to separate parameters

  adj = nx.to_scipy_sparse_matrix(G, weight='wgt')
  features = sp.lil_matrix(np.matrix(features))
  features, _ = process.preprocess_features(features)

  nb_nodes = features.shape[0]
  ft_size = features.shape[1]

  adj = process.normalize_adj(adj + sp.eye(adj.shape[0]))

  if sparse:
    sp_adj = process.sparse_mx_to_torch_sparse_tensor(adj)
  else:
    adj = (adj + sp.eye(adj.shape[0])).todense()

  features = torch.FloatTensor(features[np.newaxis])
  if not sparse:
    adj = torch.FloatTensor(adj[np.newaxis])

  model = DGI(ft_size, hid_units, nonlinearity)
  optimiser = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2_coef)

  if torch.cuda.is_available():
    print('Using CUDA')
    model.cuda()
    features = features.cuda()
    if sparse:
        sp_adj = sp_adj.cuda()
    else:
        adj = adj.cuda()

  b_xent = nn.BCEWithLogitsLoss()
  xent = nn.CrossEntropyLoss()
  cnt_wait = 0
  best = 1e9
  best_t = 0

  for epoch in range(nb_epochs):
    model.train()
    optimiser.zero_grad()
    
    idx = np.random.permutation(nb_nodes)
    shuf_fts = features[:, idx, :]

    lbl_1 = torch.ones(batch_size, nb_nodes)
    lbl_2 = torch.zeros(batch_size, nb_nodes)
    lbl = torch.cat((lbl_1, lbl_2), 1)

    if torch.cuda.is_available():
        shuf_fts = shuf_fts.cuda()
        lbl = lbl.cuda()
    
    logits = model(features, shuf_fts, sp_adj if sparse else adj, sparse, None, None, None) 
    loss = b_xent(logits, lbl)

    print('Loss:', loss)
    if loss < best:
        best = loss
        best_t = epoch
        cnt_wait = 0
    else:
        cnt_wait += 1

    if cnt_wait == patience:
        print("epochs: ", epoch)
        print('Early stopping!')
        break

    loss.backward()
    optimiser.step()
  return (((model.embed(features, sp_adj if sparse else adj, sparse, None)[0]).squeeze()).data).cpu().numpy()

[((23, 8, 23, 50), 'networkx.to_scipy_sparse_matrix', 'nx.to_scipy_sparse_matrix', (), '', True, 'import networkx as nx\n'), ((25, 16, 25, 53), 'embed_methods.dgi.utils.process.preprocess_features', 'process.preprocess_features', ({(25, 44, 25, 52): 'features'}, {}), '(features)', False, 'from embed_methods.dgi.utils import process\n'), ((37, 13, 37, 52), 'torch.FloatTensor', 'torch.FloatTensor', ({(37, 31, 37, 51): 'features[np.newaxis]'}, {}), '(features[np.newaxis])', False, 'import torch\n'), ((41, 10, 41, 47), 'embed_methods.dgi.models.DGI', 'DGI', ({(41, 14, 41, 21): 'ft_size', (41, 23, 41, 32): 'hid_units', (41, 34, 41, 46): 'nonlinearity'}, {}), '(ft_size, hid_units, nonlinearity)', False, 'from embed_methods.dgi.models import DGI, LogReg\n'), ((44, 5, 44, 30), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((53, 11, 53, 33), 'torch.nn.BCEWithLogitsLoss', 'nn.BCEWithLogitsLoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((54, 9, 54, 30), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((24, 27, 24, 46), 'numpy.matrix', 'np.matrix', ({(24, 37, 24, 45): 'features'}, {}), '(features)', True, 'import numpy as np\n'), ((33, 13, 33, 58), 'embed_methods.dgi.utils.process.sparse_mx_to_torch_sparse_tensor', 'process.sparse_mx_to_torch_sparse_tensor', ({(33, 54, 33, 57): 'adj'}, {}), '(adj)', False, 'from embed_methods.dgi.utils import process\n'), ((39, 10, 39, 44), 'torch.FloatTensor', 'torch.FloatTensor', ({(39, 28, 39, 43): 'adj[np.newaxis]'}, {}), '(adj[np.newaxis])', False, 'import torch\n'), ((63, 10, 63, 41), 'numpy.random.permutation', 'np.random.permutation', ({(63, 32, 63, 40): 'nb_nodes'}, {}), '(nb_nodes)', True, 'import numpy as np\n'), ((66, 12, 66, 44), 'torch.ones', 'torch.ones', ({(66, 23, 66, 33): 'batch_size', (66, 35, 66, 43): 'nb_nodes'}, {}), '(batch_size, nb_nodes)', False, 'import torch\n'), ((67, 12, 67, 45), 'torch.zeros', 'torch.zeros', ({(67, 24, 67, 34): 'batch_size', (67, 36, 67, 44): 'nb_nodes'}, {}), '(batch_size, nb_nodes)', False, 'import torch\n'), ((68, 10, 68, 38), 'torch.cat', 'torch.cat', ({(68, 20, 68, 34): '(lbl_1, lbl_2)', (68, 36, 68, 37): '1'}, {}), '((lbl_1, lbl_2), 1)', False, 'import torch\n'), ((70, 7, 70, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((30, 36, 30, 56), 'scipy.sparse.eye', 'sp.eye', ({(30, 43, 30, 55): 'adj.shape[0]'}, {}), '(adj.shape[0])', True, 'import scipy.sparse as sp\n'), ((35, 17, 35, 37), 'scipy.sparse.eye', 'sp.eye', ({(35, 24, 35, 36): 'adj.shape[0]'}, {}), '(adj.shape[0])', True, 'import scipy.sparse as sp\n')]