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()

[((438, 480), 'networkx.to_scipy_sparse_matrix', 'nx.to_scipy_sparse_matrix', (['G'], {'weight': '"""wgt"""'}), "(G, weight='wgt')\n", (463, 480), True, 'import networkx as nx\n'), ((545, 582), 'embed_methods.dgi.utils.process.preprocess_features', 'process.preprocess_features', (['features'], {}), '(features)\n', (572, 582), False, 'from embed_methods.dgi.utils import process\n'), ((848, 887), 'torch.FloatTensor', 'torch.FloatTensor', (['features[np.newaxis]'], {}), '(features[np.newaxis])\n', (865, 887), False, 'import torch\n'), ((961, 998), 'embed_methods.dgi.models.DGI', 'DGI', (['ft_size', 'hid_units', 'nonlinearity'], {}), '(ft_size, hid_units, nonlinearity)\n', (964, 998), False, 'from embed_methods.dgi.models import DGI, LogReg\n'), ((1085, 1110), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1108, 1110), False, 'import torch\n'), ((1277, 1299), 'torch.nn.BCEWithLogitsLoss', 'nn.BCEWithLogitsLoss', ([], {}), '()\n', (1297, 1299), True, 'import torch.nn as nn\n'), ((1309, 1330), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), '()\n', (1328, 1330), True, 'import torch.nn as nn\n'), ((508, 527), 'numpy.matrix', 'np.matrix', (['features'], {}), '(features)\n', (517, 527), True, 'import numpy as np\n'), ((731, 776), 'embed_methods.dgi.utils.process.sparse_mx_to_torch_sparse_tensor', 'process.sparse_mx_to_torch_sparse_tensor', (['adj'], {}), '(adj)\n', (771, 776), False, 'from embed_methods.dgi.utils import process\n'), ((915, 949), 'torch.FloatTensor', 'torch.FloatTensor', (['adj[np.newaxis]'], {}), '(adj[np.newaxis])\n', (932, 949), False, 'import torch\n'), ((1465, 1496), 'numpy.random.permutation', 'np.random.permutation', (['nb_nodes'], {}), '(nb_nodes)\n', (1486, 1496), True, 'import numpy as np\n'), ((1545, 1577), 'torch.ones', 'torch.ones', (['batch_size', 'nb_nodes'], {}), '(batch_size, nb_nodes)\n', (1555, 1577), False, 'import torch\n'), ((1590, 1623), 'torch.zeros', 'torch.zeros', (['batch_size', 'nb_nodes'], {}), '(batch_size, nb_nodes)\n', (1601, 1623), False, 'import torch\n'), ((1634, 1662), 'torch.cat', 'torch.cat', (['(lbl_1, lbl_2)', '(1)'], {}), '((lbl_1, lbl_2), 1)\n', (1643, 1662), False, 'import torch\n'), ((1671, 1696), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1694, 1696), False, 'import torch\n'), ((682, 702), 'scipy.sparse.eye', 'sp.eye', (['adj.shape[0]'], {}), '(adj.shape[0])\n', (688, 702), True, 'import scipy.sparse as sp\n'), ((802, 822), 'scipy.sparse.eye', 'sp.eye', (['adj.shape[0]'], {}), '(adj.shape[0])\n', (808, 822), True, 'import scipy.sparse as sp\n')]