import _init_paths
import utils
import graph
import networkx as nx
import torch
try:
    import lib.graph as gwGraph
    from lib.ot_distances import Fused_Gromov_Wasserstein_distance
except:
    print('NOTE: GW library not found. Clone the GW repo as in README if needed.')
import numpy as np
import pickle
from tqdm import tqdm
from torch.utils.data.dataloader import DataLoader
from torch.utils.data import Dataset
from sklearn.model_selection import GridSearchCV
import got_stochastic as st
import torch.nn as nn
import sys
import runGraph
import os
import math
import netlsd
import numpy.linalg as lg

from sklearn.pipeline import Pipeline
import grakel
from sklearn.svm import SVC
import time

import pdb
from joblib import Parallel, delayed 
torch.set_default_tensor_type('torch.DoubleTensor')

from searchGraph import * 
from variation_coarsen import * 
from otc import *    

os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3,4'

# datasets = ['MSRC_21C', 'BZR', 'MUTAG', 'SYNTHETIC', 'MSRC_21', 'DHFR', 'COX2_MD', 'Synthie']
datasets = ['BZR', 'MUTAG', 'Synthie']
algorithms = ['variation_neighborhood', 'variation_edges', 'affinity_GS', 'algebraic_JC', 'heavy_edge', 'otc', 'copt', 'ergot']

results_np = []
      
tgt_n = None
for dataset_name in tqdm(datasets):
      args = utils.parse_args()
      args.dataset_name = dataset_name 
      args.n_epochs = 400
      for algo in tqdm(algorithms[:-2]):
            torch.cuda.empty_cache()
            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)
            args.algo = algo
            args.compress_fac = 4
            print(f'>>> Currently computing {algo} on {dataset_name} >>>')
            filter_type = 'L'
            sketch_path = 'sketch/{}_data_sketch{}_{}_{}_{}.pt'.format(dataset_name, args.compress_fac, args.hike_interval, filter_type, algo)
            if algo in ['variation_neighborhood', 'variation_edges', 'affinity_GS', 'algebraic_JC', 'heavy_edge']:
                        
                  # if os.path.exists(sketch_path):
                  #       print('Using cached data!')
                  #       data_sketch = torch.load(sketch_path)
                  #       node_labels, data_sketch = data_sketch['sketch_P'], data_sketch['data_sketch']
                  # else:
                  if not os.path.exists(sketch_path):
                        # os.remove(sketch_path)
                        print('Creating and caching data!')
                        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)
                        
                        if not os.path.exists(sketch_path):
                              continue 
            elif algo == 'otc':
      
                  # if os.path.exists(sketch_path):
                  #       print('Using cached data!')
                  #       data_sketch = torch.load(sketch_path)
                  #       node_labels, data_sketch = data_sketch['sketch_P'], data_sketch['data_sketch']
                  # else:
                  # if os.path.exists(sketch_path):
                  #       os.remove(sketch_path)
                  if not os.path.exists(sketch_path):
                        print('Creating and caching data!')
                        data_sketch, sketch_P = otc_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)
                        if not os.path.exists(sketch_path):
                              continue 
                  
            # else:
                  # if os.path.exists(sketch_path):
                  #       os.remove(sketch_path)
            
            try:
                  # args.C = 3
            # print(os.path.exists(sketch_path))
                  args.compress_fac = 4
      # args.grid_search = True
                  accs = test_sketch_real_data(args)
            except Exception as e:
                  print(e)
                  continue 
            results_np.append([dataset_name, algo] + accs.tolist())
            
            np.save('sketching_results_first.npy', results_np)
            