import numpy as np 
import numpy.random as npr

import torch
import cvxpy as cp
import networkx as nx 
import utils
import math 
from tqdm import tqdm 

def algebraic_distance_dense(G, params={}):
      '''
      takes: graph G, computational parameters
      returns:
      a distance dictionary, d[node1][node2]  giving the distance between the nodes
      ref:
                  RELAXATION-BASED COARSENING AND
                  MULTISCALE GRAPH ORGANIZATION
                  DORIT RON, ILYA SAFRO, AND ACHI BRANDT
      '''

      metric              = params.get('metric', 'Linfinity')
      num_relaxations_r   = params.get('num_relaxations_r', 10)
      num_test_vectors_K  = params.get('num_test_vectors_K', 20)
      lazy_walk_param_w   = params.get('lazy_walk_param_w', 0.3)

      if metric != 'Linfinity':
            raise Exception('Metric other than Linifinity not implemented')

      H = nx.convert_node_labels_to_integers(G, label_attribute='orig_label')
      all_nodes = H.nodes()  #future: use this for sorting

      distance = {}
      for node1 in all_nodes:
            distance[node1] = {}
            for node2 in H.neighbors(node1):
                  if node1 < node2: #save time
                        distance[node1][node2] = -np.inf

      #wishlist: sparse matrices
      #the wrong laplacian (uses degrees) LAP      = nx.laplacian_matrix(H)
      #should consider degree? diag_vec     = np.array([H.node[u].get('weight', 1) for u in H])
      #diag_vec     = np.array([H.degree(u) for u in H])
      diag_vec = []
      for u in all_nodes:
            # print(H.edges(u))
            val = 0. + len(H.edges(u))
            if val == 0.:
                  val = 1.
            diag_vec.append(val)
      diag_vec     = np.array(diag_vec)
      diag_vec_inv = 1./diag_vec  #[1./val for val in diag_vec]
      DIAG     = np.diag(diag_vec)
      DIAGinv  = np.diag(diag_vec_inv)
      ADJ      = nx.adj_matrix(G, nodelist=all_nodes)
      w_times_Dinv_times_LAP = lazy_walk_param_w * np.dot(DIAGinv,DIAG-ADJ)

      for t in range(num_test_vectors_K):
            x = npr.rand(H.number_of_nodes(),1) - 0.5
            x = x / x.sum()

            for iteration in range(num_relaxations_r):
                  x = (1-lazy_walk_param_w)*x + np.dot(w_times_Dinv_times_LAP, x)

            #maximize over the trial vectors: d(i,j) = max_{t=1..K} |x_t(i) - x_t(j)|
            for node1 in all_nodes:
                  for node2 in H.neighbors(node1):
                        dis = np.abs((x[node1]-x[node2])[0,0])
                        if node1 < node2 and dis > distance[node1][node2]: #to save time, compute just the upper triangle of the matrix
                              distance[node1][node2] = dis


      #generate the distance dictionary in the original node labels, and including the diagonal and lower triangle
      ret = {}
      for node1 in G:
            ret[node1] = {node1:0.}
      for u in H:
            node1 = H.nodes[u]['orig_label']
            ret[node1] = {}
            for v in H.neighbors(u):
                  node2 = H.nodes[v]['orig_label']
                  if u < v:
                        d = distance[u][v]
                  elif v < u:
                        d = distance[v][u]
                  else:
                        d = 0.
                  ret[node1][node2] = d
                  ret[node2][node1] = d
      return ret

def select_coarse_nodes(d, g, n, Q=0.5, nu=2):
      H = nx.convert_node_labels_to_integers(g, label_attribute='orig_label')
      F = set(d.keys())
      theta = {}
      v = [1] * len(F)
      for i,u in enumerate(F):
            theta[u] = v[i] + sum([1/d[u][s]/sum([1/d[s][t] for t in H.neighbors(s)]) for s in H.neighbors(u)])
            
      avg_theta = sum(list(theta.values()))/len(list(theta.values()))
      C = set([k for k, t in theta.items() if t > nu * avg_theta])
      F = set(d.keys()) - set(C)
      theta_f = dict(sorted({k:t for k,t in theta.items() if k in F}.items(), key=lambda item:item[1]), reverse=True)
      for i, k in enumerate(F):
            neighbor = H.neighbors(k)
            n_c = set(neighbor) & set(C)
            if (sum([1/d[k][t] for t in neighbor]) <= Q * sum([1/d[k][t] for t in n_c])) or \
                  (len(list(neighbor)) <= Q * len(list(n_c))):
                  F = F - set([k])
                  C = C | set([k])
                  if len(list(C)) >= n:
                        break
      return list(C)
      

def coarsen_graph(g, n, Q=0.5, v=2):
      d = algebraic_distance_dense(g)
      C = select_coarse_nodes(d, g, n, Q, v)
      return C
      
def algebra_sketch(dataset, node_labels, args, tgt_n):
      min_label = min([min(l) for l in node_labels])
      max_label = max([max(l) for l in node_labels])
      node_labels -= min_label
      n_labels = int(max_label - min_label) + 1
      data_sketch = []
      P_sketch = []
      for i, data in enumerate(tqdm(dataset)):
            g = utils.lap_to_graph(data)
            cur_labels = node_labels[i]
            # print(cur_labels, cur_labels.shape, data.shape)
            edge_labels = [(cur_labels[e[0]], cur_labels[e[1]]) for e in g.edges]
            c = np.array([0 if x[0] == x[1] else 1 for x in edge_labels]) * 0.01 + 0.01 
            args.m = len(data)
            args.n = math.ceil(args.m/args.compress_fac) if tgt_n is None else tgt_n
            if args.n >= args.m:
                  continue 
            top_n_idx = coarsen_graph(g, args.n)
            print(len(g.nodes()), args.m, args.n)
            sketched_graph = nx.induced_subgraph(g, top_n_idx)
            sketch_labels = node_labels[i][top_n_idx].astype(int)
            prob_labels = np.zeros((n_labels, args.n))
            prob_labels[sketch_labels, np.arange(args.n)] = 1
            L = utils.graph_to_lap(sketched_graph)
            L_ordering = L.diag().argsort(dim=0)
            sketch_labels = sketch_labels[L_ordering]
            
            ones = torch.ones(args.n, args.n).triu()
            L = utils.canonicalize_mx(L)
            L = L[ones > 0]
            
            data_sketch.append(L)
            P_sketch.append(torch.from_numpy(prob_labels).t())
            
      return data_sketch, P_sketch

if __name__ == '__main__':
      args = utils.parse_args()
      dataset_name = args.dataset_name
      enz_data = torch.load('data/{}_lap.pt'.format(dataset_name))
      args.compress_fac = args.compress_fac if args.compress_fac > 0  else 2
      print('dataset: {} compress fac {}'.format( dataset_name, args.compress_fac))
      dataset = enz_data['lap']
      target = np.array(enz_data['target'])
      node_labels = enz_data['labels']
      
      algo = args.algo
      method2tgt_n = {'PROTEINS': 13, 'BZR_MD': 7, 'MSRC_9':7, 'ENZYMES':10, 'BZR': 7, 'MUTAG':7, 'SYNTHETIC':50}
      tgt_n = method2tgt_n[dataset_name]
      sketch_path = '{}_data_sketch{}tgt_{}_{}.pt'.format(dataset_name, tgt_n, args.hike_interval, algo)
      
      data_sketch, sketch_P = algebra_sketch(dataset, node_labels, args, tgt_n)
      node_labels = sketch_P
      torch.save({'data_sketch':data_sketch, 'sketch_P':sketch_P, 'compress_fac':args.compress_fac} , sketch_path)
      # data_sketch = torch.load(sketch_path)
      # node_labels, data_sketch = data_sketch['sketch_P'], data_sketch['data_sketch']
      # print(node_labels[0].shape)
            