import networkx as nx
import numpy as np
import scipy as sp
import scipy.cluster.vq as vq
import matplotlib.pyplot as plt
import math
import random
import operator
from networkx import graph_atlas_g

################################################################################
import torch
import torch_geometric
import torch.optim as optim
from torch_geometric.nn import GCNConv
from torch_geometric.utils import from_networkx
from torch_geometric.datasets import Planetoid
from torch_geometric.utils import to_networkx

################################################################################
################################################################################
################  Part 1: GCN Vertex Classification Analysis  ##################
################################################################################
################################################################################
# Use this model as a starting point. You'll primarily be modifying the 
# variables associated with the dataset (x, num_features, train_mask, test_mask)
class GCN(torch.nn.Module):
  def __init__(self, data):
    super(GCN, self).__init__()
    self.conv1 = torch_geometric.nn.GCNConv(data.num_features, 64)
    self.conv2 = torch_geometric.nn.GCNConv(64, 64)
    self.conv3 = torch_geometric.nn.GCNConv(64, 64)
    self.classifier = torch.nn.Linear(64, data.num_classes)
  
  def forward(self, x, edge_index):
    h = self.conv1(x, edge_index)
    h = h.tanh()
    h = self.conv2(h, edge_index)
    h = h.tanh()
    h = self.conv3(h, edge_index)
    h = h.tanh()
    out = self.classifier(h)
    return out, h

# training function
def train(model, criterion, optimizer, data):
  optimizer.zero_grad()
  out, h = model(data.x, data.edge_index)
  loss = criterion(out[data.train_mask], data.y[data.train_mask])
  loss.backward()
  optimizer.step()
  return loss, h

# test function
@torch.no_grad()
def test(model, data):
  model.eval()
  logits, _ = model(data.x, data.edge_index)
  pred = logits[data.test_mask].max(1)[1]
  acc = pred.eq(data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
  return acc

################################################################################
# Load up the Cora dataset from Planetoid


# Build the different feature sets as described in the HW4 document


# Perform the training and test loop, track the maximum accuracies per features
max_acc_lp = 0.0
max_acc_base = 0.0
max_acc_struct = 0.0
max_acc_cent = 0.0
max_acc_struct_cent = 0.0
max_acc_base_struct_cent = 0.0


# Output our results
print("Max accuracy (label prop):", max_acc_lp)
print("Max accuracy (baseline):", max_acc_base)
print("Max accuracy (structure only):", max_acc_struct)
print("Max accuracy (centrality only):", max_acc_cent)
print("Max accuracy (structure+centrality):", max_acc_struct_cent)
print("Max accuracy (baseline+structure+centrality):", max_acc_base_struct_cent)


################################################################################
################################################################################
#########################  Part 2: Null Model Usage   ##########################
################################################################################
################################################################################
# Read in graph and extract the n=order, m=size, and degree distribution
G = nx.read_edgelist("dolphins.data", comments="%")


################################################################################
# Generate 5 each of:
#   a.) Erdos-Renyi Graphs (using n, m)
#   b.) Chung-Lu Graphs (using degree distribution)
#   c.) Barabasi-Albert Graphs (using n)
#   d.) Watts-Strogatz Graphs (using n, m, p=0.25)


################################################################################
# Count the average occurance of all 3- and 4-vertex connected subgraphs in the
# above graphs as well as the real world network

# Generate all 3- and 4-vertex connected subgraphs
all_subgraphs = list()
all_subgraphs.append([g for g in graph_atlas_g() if len(g.nodes())==3 and len(list(nx.connected_components(g)))==1])
all_subgraphs.append([g for g in graph_atlas_g() if len(g.nodes())==4 and len(list(nx.connected_components(g)))==1])


# Initialize the arrays for tracking average counts
counts_real = list()
counts_er = list()
counts_cl = list()
counts_ba = list()
counts_ws = list()
for i in range(0, len(all_subgraphs)):
  counts_real.append(0.0)
  counts_er.append(0.0)
  counts_cl.append(0.0)
  counts_ba.append(0.0)
  counts_ws.append(0.0)


# Get the counts for each subgraph in all_subgraphs, determine average
# Note: it might take a while for 4-vertex subgraphs, test with a single graph


################################################################################
# Output the average results for each subgraph

for i in range(0, len(all_subgraphs)):
  print("Subgraph: ", i, "Real:", counts_real[i], "Erdos-Renyi:", counts_er[i], "Chung-Lu:", counts_cl[i], "Barabasi-Albert:", counts_ba[i], "Watts-Strogatz:", counts_ws[i])