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 

def Phi_gradient(eps, t, zeta, rho, lam):
      # print(type(t), type(zeta))
      # large_t = t[t > -zeta]
      # small_t = t[t <= -zeta]
      # sum_1 = -1/2/lam * (eps * torch.square(small_t + zeta)).sum()
      
      # sum_2 = - (t * rho).sum()
      # sum_3 = -zeta 
      
      eps_grad = np.zeros(len(eps))
      t_grad = - rho
      large_idx = t > -zeta
      small_idx = t <= -zeta
      eps_grad[small_idx] = -np.square(t[small_idx] + zeta)/2/lam 
      t_grad[large_idx] = -rho[large_idx]
      t_grad[small_idx] = -1/2/lam * (2 * np.multiply(eps[small_idx], t[small_idx] + zeta))
      zeta_grad = -1/lam * (t[small_idx].sum() + small_idx.sum() * zeta) - 1
      
      return eps_grad, t_grad, zeta_grad
      # print(sum_1, sum_2, sum_3)
      # return sum_1 + sum_2 + sum_3

def proj_eps(eps, k, tol=1e-3):
      eps = np.clip(eps, 0, 1)
      if eps.sum() <= k:
            return eps 
      n = len(eps)
      l = (eps.sum() - k)/n
      h = eps.max() - k/n 
      while (l < h) or (np.abs(l-h) < 1e-10):
            mid = (l + h)/2 
            eps_hat = eps - mid 
            eps_hat = np.clip(eps_hat, 0, 1)
            if np.abs(eps_hat.sum() - k) < tol:
                  return eps_hat 
            elif eps_hat.sum() > k:
                  l = mid + 0.
            else:
                  h = mid + 0. 
            
            
def proj_t(t, F, c):
      x = cp.Variable(len(t))
      loss = cp.sum_squares(t - x)
      prob = cp.Problem(cp.Minimize(loss),
                        [F @ x <= c, F @ x >= -c])
      prob.solve()
      
      return x.value
      
# def set_grad(var):
#           def hook(grad):
#         var.grad = grad
#     return hook


def otc_solve(g, c, ori_dim, compre_dim, T=25, lam=1):
      F = nx.incidence_matrix(g).todense().T
      deg = np.array([val for _, val in g.degree()])
      rho = deg / deg.sum() 
      
      eps = np.ones(ori_dim)/ori_dim * compre_dim
      # eps.requires_grad = True
      t = np.zeros(ori_dim)
      # t.requires_grad = True
      zeta = 0
      # zeta.requires_grad = True
      alpha, beta, gamma = 0.1, 0.1, 0.1
      ret_eps_hat = np.zeros(ori_dim)
      cum_alpha = 0
      for i in range(T):
            eps_grad, t_grad, zeta_grad = Phi_gradient(eps, t, zeta, rho, lam)
            # phi.backward()
            
            eps_hat = proj_eps(eps - alpha * eps_grad, k=compre_dim)
            # eps_hat.requires_grad = True
            t_hat = proj_t(t + beta * t_grad, F, c)
            # t_hat.requires_grad = True
            zeta_hat = zeta + gamma * zeta_grad
            # zeta_hat.requires_grad = True

            eps_hat_grad, t_hat_grad, zeta_hat_grad = Phi_gradient(eps_hat, t_hat, zeta_hat, rho, lam)
            # phi_hat.backward()
            eps = proj_eps(eps - alpha * eps_hat_grad, k=compre_dim)
            # eps.requires_grad = True
            t = proj_t(t + beta * t_hat_grad, F, c)
            # t.requires_grad = True
            zeta = zeta + gamma * zeta_hat_grad
            # zeta.requires_grad = True
            
            ret_eps_hat += alpha * eps_hat
            cum_alpha += alpha
      return ret_eps_hat/cum_alpha


def otc_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 
            eps = otc_solve(g, c, args.m, args.n)
            # print(len(g.nodes()), args.m, args.n)
            top_n_idx = np.argpartition(eps, args.n)[-args.n:]
            sketched_graph = nx.induced_subgraph(g, top_n_idx)
            sketch_labels = node_labels[i][top_n_idx].astype(int)
            # prob_labels = torch.zeros(n_labels, args.n)
            # cur_labels = torch.from_numpy(node_labels[i]).to(torch.int64).unsqueeze(-1).expand(-1, args.n)
            # print(n_labels, node_labels[i].shape, cur_labels)
            # print(top_n_idx)
            # prob_labels[cur_labels] = 1
            # print(prob_labels)
            # break 
            prob_labels = np.zeros((n_labels, args.n))
            prob_labels[sketch_labels, np.arange(args.n)] = 1
            # print(prob_labels.T)
            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']
      
      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 = 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)
      # data_sketch = torch.load(sketch_path)
      # node_labels, data_sketch = data_sketch['sketch_P'], data_sketch['data_sketch']
      # print(node_labels[0].shape)
            