import numpy as np 
import torch
# from searchGraph import sketch_graph 
import cvxpy as cp
import networkx as nx 
import utils
import math 
from tqdm import tqdm 
import pygsp as gsp 
from graph_coarsening.coarsening_utils import *
import graph_coarsening.graph_utils

def var_sketch(dataset, node_labels, args, tgt_n):
      # print(node_labels)
      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)
            G = gsp.graphs.Graph.from_networkx(g)
            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 
            r = 1 - args.n/args.m
            k = 5  
            method = args.algo 
            try:
                  C, Gc, Call, Gall, rows_remain = coarsen(G, K=k, n_target=args.n, method=method, max_levels=100)
            except:
                  continue 
            if not (len(rows_remain) == args.n):
                  continue 
            top_n_idx = np.array(rows_remain)
            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']
      np.random.seed(args.seed)
      filter_type = 'L'
      algo = args.algo
      method2tgt_n = {'PROTEINS':13, 'BZR_MD':7, 'MSRC_9':20, 'MSRC_21C':20,  'ENZYMES':10, 'BZR': 7, 'MUTAG':7, 'SYNTHETIC':50, 'MSRC_21':40, 'DHFR':30, 'COX2_MD':13, 'Synthie':50}
      tgt_n = method2tgt_n[dataset_name]
      # tgt_n = None 
      if tgt_n is not None:
            sketch_path = 'retrieval/{}_data_sketch{}tgt_{}_{}.pt'.format(dataset_name, tgt_n, args.hike_interval, algo)
      else:
            sketch_path = 'sketch/{}_data_sketch{}_{}_{}_{}.pt'.format(dataset_name, args.compress_fac, args.hike_interval, filter_type, algo)
      
      
      data_sketch, sketch_P = var_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)
            