"""
"""

import numpy as np

import torch
from torch import optim
import torch.nn.functional as F

from scipy.spatial import distance
import scipy.sparse as sp
import networkx as nx

import auto_encoders.gae.utils as gae_util
import auto_encoders.gae.optimizer as gae_optimizer

import auto_encoders.model as model
import info_log

def graph_AE_handler(X_embed, CCC_graph, args, param):
    info_log.print('--------> Starting Graph AE ...')

    use_GAT = args.graph_AE_use_GAT
    learning_rate = args.graph_AE_learning_rate
    total_epoch = args.graph_AE_epoch
    embedding_size = args.graph_AE_embedding_size

    # Prepare matrices
    if use_GAT:
        X_embed_normalized = normalize_features_dense(X_embed)
        X_embed_normalized = torch.from_numpy(X_embed_normalized).type(torch.FloatTensor).to(param['device'])

        CCC_graph_edge_index = convert_adj_to_edge_index(CCC_graph)
        CCC_graph_edge_index = torch.from_numpy(CCC_graph_edge_index).type(torch.LongTensor).to(param['device'])

        CCC_graph = torch.from_numpy(CCC_graph).type(torch.FloatTensor).to(param['device'])
    else:
        adj, adj_train, edgeList = feature2adj(X_embed)
        adj_norm = gae_util.preprocess_graph(adj)
        adj_label = (adj_train + sp.eye(adj_train.shape[0])).toarray()

        zDiscret = X_embed > np.mean(X_embed, axis=0)
        zDiscret = 1.0 * zDiscret
        X_embed_normalized = torch.from_numpy(zDiscret).type(torch.FloatTensor).to(param['device'])
        CCC_graph_edge_index = adj_norm.to(param['device'])
        CCC_graph = torch.from_numpy(adj_label).type(torch.FloatTensor).to(param['device'])

        pos_weight = float(adj.shape[0] * adj.shape[0] - adj.sum()) / adj.sum()
        norm = adj.shape[0] * adj.shape[0] / float((adj.shape[0] * adj.shape[0] - adj.sum()) * 2)

    graph_AE = model.Graph_AE(X_embed.shape[1], embedding_size).to(param['device'])
    optimizer = optim.Adam(graph_AE.parameters(), lr=learning_rate)

    for epoch in range(total_epoch):
        graph_AE.train()
        optimizer.zero_grad()

        embed, gae_info, recon_graph = graph_AE(X_embed_normalized, CCC_graph_edge_index, use_GAT=use_GAT)

        if use_GAT:
            loss = loss_function(preds = recon_graph,
                                labels = CCC_graph)
        else:
            loss = gae_optimizer.loss_function(preds = recon_graph,
                                            labels = CCC_graph,
                                            mu=gae_info[0], logvar=gae_info[1],
                                            n_nodes = X_embed.shape[0],
                                            norm = norm, 
                                            pos_weight = pos_weight)
        
        # Backprop and Update
        loss.backward()
        cur_loss = loss.item()
        optimizer.step()

        info_log.interval_print(f"----------------> Epoch: {epoch+1}/{total_epoch}, Current loss: {cur_loss:.4f}", epoch=epoch, total_epoch=total_epoch)

    return embed.detach().cpu().numpy(), recon_graph.detach().cpu().numpy(), edgeList, adj # edgeList added just for benchmark testing

def loss_function(preds, labels):
    return F.binary_cross_entropy_with_logits(preds, labels)

def normalize_features_dense(node_features_dense):
    assert isinstance(node_features_dense, np.ndarray), f'Expected np matrix got {type(node_features_dense)}.'

    # The goal is to make feature vectors normalized (sum equals 1), but since some feature vectors are all 0s
    # in those cases we'd have division by 0 so I set the min value (via np.clip) to 1.
    # Note: 1 is a neutral element for division i.e. it won't modify the feature vector
    return node_features_dense / np.clip(node_features_dense.sum(1,keepdims=True), a_min=1, a_max=None)

def convert_adj_to_edge_index(adjacency_matrix):
    """
    """
    assert isinstance(adjacency_matrix, np.ndarray), f'Expected NumPy array got {type(adjacency_matrix)}.'
    height, width = adjacency_matrix.shape
    assert height == width, f'Expected square shape got = {adjacency_matrix.shape}.'

    # If there are infs that means we have a connectivity mask and 0s are where the edges in connectivity mask are,
    # otherwise we have an adjacency matrix and 1s symbolize the presence of edges.
    # active_value = 0 if np.isinf(adjacency_matrix).any() else 1

    edge_index = []
    for src_node_id in range(height):
        for trg_nod_id in range(width):
            if adjacency_matrix[src_node_id, trg_nod_id] > 0:
                edge_index.append([src_node_id, trg_nod_id])

    return np.asarray(edge_index).transpose()  # change shape from (N,2) -> (2,N)

def feature2adj(X_embed):
    edgeList = calculateKNNgraphDistanceMatrixStatsSingleThread(X_embed)
    graphdict = edgeList2edgeDict(edgeList, X_embed.shape[0])
    adj = nx.adjacency_matrix(nx.from_dict_of_lists(graphdict))
    adj_orig = adj
    adj_orig = adj_orig - sp.dia_matrix((adj_orig.diagonal()[np.newaxis, :], [0]), shape=adj_orig.shape)
    adj_orig.eliminate_zeros()
    adj_train, train_edges, val_edges, val_edges_false, test_edges, test_edges_false = gae_util.mask_test_edges(adj)
    adj = adj_train

    return adj, adj_train, edgeList


def calculateKNNgraphDistanceMatrixStatsSingleThread(featureMatrix, distanceType='euclidean', k=10):
    r"""
    Thresholdgraph: KNN Graph with stats one-std based methods, SingleThread version
    """       

    edgeList=[]

    for i in np.arange(featureMatrix.shape[0]):
        tmp=featureMatrix[i,:].reshape(1,-1)
        distMat = distance.cdist(tmp, featureMatrix, distanceType)
        res = distMat.argsort()[:k+1]
        tmpdist = distMat[0, res[0][1:k+1]]
        boundary = np.mean(tmpdist) + np.std(tmpdist)
        for j in np.arange(1, k+1):
            # TODO: check, only exclude large outliners
            # if (distMat[0,res[0][j]]<=mean+std) and (distMat[0,res[0][j]]>=mean-std):
            if distMat[0,res[0][j]]<=boundary:
                weight = 1.0
            else:
                weight = 0.0
            edgeList.append((i, res[0][j], weight))
    
    return edgeList

def edgeList2edgeDict(edgeList, nodesize):
    graphdict={}
    tdict={}

    for edge in edgeList:
        end1 = edge[0]
        end2 = edge[1]
        tdict[end1]=""
        tdict[end2]=""
        if end1 in graphdict:
            tmplist = graphdict[end1]
        else:
            tmplist = []
        tmplist.append(end2)
        graphdict[end1]= tmplist

    #check and get full matrix
    for i in range(nodesize):
        if i not in tdict:
            graphdict[i]=[]

    return graphdict
"""
"""

import numpy as np

import torch
from torch import optim
import torch.nn.functional as F

from scipy.spatial import distance
import scipy.sparse as sp
import networkx as nx

import auto_encoders.gae.utils as gae_util
import auto_encoders.gae.optimizer as gae_optimizer

import auto_encoders.model as model
import info_log

def graph_AE_handler(X_embed, CCC_graph, args, param):
    info_log.print('--------> Starting Graph AE ...')

    use_GAT = args.graph_AE_use_GAT
    learning_rate = args.graph_AE_learning_rate
    total_epoch = args.graph_AE_epoch
    embedding_size = args.graph_AE_embedding_size

    # Prepare matrices
    if use_GAT:
        X_embed_normalized = normalize_features_dense(X_embed)
        X_embed_normalized = torch.from_numpy(X_embed_normalized).type(torch.FloatTensor).to(param['device'])

        CCC_graph_edge_index = convert_adj_to_edge_index(CCC_graph)
        CCC_graph_edge_index = torch.from_numpy(CCC_graph_edge_index).type(torch.LongTensor).to(param['device'])

        CCC_graph = torch.from_numpy(CCC_graph).type(torch.FloatTensor).to(param['device'])
    else:
        adj, adj_train, edgeList = feature2adj(X_embed)
        adj_norm = gae_util.preprocess_graph(adj)
        adj_label = (adj_train + sp.eye(adj_train.shape[0])).toarray()

        zDiscret = X_embed > np.mean(X_embed, axis=0)
        zDiscret = 1.0 * zDiscret
        X_embed_normalized = torch.from_numpy(zDiscret).type(torch.FloatTensor).to(param['device'])
        CCC_graph_edge_index = adj_norm.to(param['device'])
        CCC_graph = torch.from_numpy(adj_label).type(torch.FloatTensor).to(param['device'])

        pos_weight = float(adj.shape[0] * adj.shape[0] - adj.sum()) / adj.sum()
        norm = adj.shape[0] * adj.shape[0] / float((adj.shape[0] * adj.shape[0] - adj.sum()) * 2)

    graph_AE = model.Graph_AE(X_embed.shape[1], embedding_size).to(param['device'])
    optimizer = optim.Adam(graph_AE.parameters(), lr=learning_rate)

    for epoch in range(total_epoch):
        graph_AE.train()
        optimizer.zero_grad()

        embed, gae_info, recon_graph = graph_AE(X_embed_normalized, CCC_graph_edge_index, use_GAT=use_GAT)

        if use_GAT:
            loss = loss_function(preds = recon_graph,
                                labels = CCC_graph)
        else:
            loss = gae_optimizer.loss_function(preds = recon_graph,
                                            labels = CCC_graph,
                                            mu=gae_info[0], logvar=gae_info[1],
                                            n_nodes = X_embed.shape[0],
                                            norm = norm, 
                                            pos_weight = pos_weight)
        
        # Backprop and Update
        loss.backward()
        cur_loss = loss.item()
        optimizer.step()

        info_log.interval_print(f"----------------> Epoch: {epoch+1}/{total_epoch}, Current loss: {cur_loss:.4f}", epoch=epoch, total_epoch=total_epoch)

    return embed.detach().cpu().numpy(), recon_graph.detach().cpu().numpy(), edgeList, adj # edgeList added just for benchmark testing

def loss_function(preds, labels):
    return F.binary_cross_entropy_with_logits(preds, labels)

def normalize_features_dense(node_features_dense):
    assert isinstance(node_features_dense, np.ndarray), f'Expected np matrix got {type(node_features_dense)}.'

    # The goal is to make feature vectors normalized (sum equals 1), but since some feature vectors are all 0s
    # in those cases we'd have division by 0 so I set the min value (via np.clip) to 1.
    # Note: 1 is a neutral element for division i.e. it won't modify the feature vector
    return node_features_dense / np.clip(node_features_dense.sum(1,keepdims=True), a_min=1, a_max=None)

def convert_adj_to_edge_index(adjacency_matrix):
    """
    """
    assert isinstance(adjacency_matrix, np.ndarray), f'Expected NumPy array got {type(adjacency_matrix)}.'
    height, width = adjacency_matrix.shape
    assert height == width, f'Expected square shape got = {adjacency_matrix.shape}.'

    # If there are infs that means we have a connectivity mask and 0s are where the edges in connectivity mask are,
    # otherwise we have an adjacency matrix and 1s symbolize the presence of edges.
    # active_value = 0 if np.isinf(adjacency_matrix).any() else 1

    edge_index = []
    for src_node_id in range(height):
        for trg_nod_id in range(width):
            if adjacency_matrix[src_node_id, trg_nod_id] > 0:
                edge_index.append([src_node_id, trg_nod_id])

    return np.asarray(edge_index).transpose()  # change shape from (N,2) -> (2,N)

def feature2adj(X_embed):
    edgeList = calculateKNNgraphDistanceMatrixStatsSingleThread(X_embed)
    graphdict = edgeList2edgeDict(edgeList, X_embed.shape[0])
    adj = nx.adjacency_matrix(nx.from_dict_of_lists(graphdict))
    adj_orig = adj
    adj_orig = adj_orig - sp.dia_matrix((adj_orig.diagonal()[np.newaxis, :], [0]), shape=adj_orig.shape)
    adj_orig.eliminate_zeros()
    adj_train, train_edges, val_edges, val_edges_false, test_edges, test_edges_false = gae_util.mask_test_edges(adj)
    adj = adj_train

    return adj, adj_train, edgeList


def calculateKNNgraphDistanceMatrixStatsSingleThread(featureMatrix, distanceType='euclidean', k=10):
    r"""
    Thresholdgraph: KNN Graph with stats one-std based methods, SingleThread version
    """       

    edgeList=[]

    for i in np.arange(featureMatrix.shape[0]):
        tmp=featureMatrix[i,:].reshape(1,-1)
        distMat = distance.cdist(tmp, featureMatrix, distanceType)
        res = distMat.argsort()[:k+1]
        tmpdist = distMat[0, res[0][1:k+1]]
        boundary = np.mean(tmpdist) + np.std(tmpdist)
        for j in np.arange(1, k+1):
            # TODO: check, only exclude large outliners
            # if (distMat[0,res[0][j]]<=mean+std) and (distMat[0,res[0][j]]>=mean-std):
            if distMat[0,res[0][j]]<=boundary:
                weight = 1.0
            else:
                weight = 0.0
            edgeList.append((i, res[0][j], weight))
    
    return edgeList

def edgeList2edgeDict(edgeList, nodesize):
    graphdict={}
    tdict={}

    for edge in edgeList:
        end1 = edge[0]
        end2 = edge[1]
        tdict[end1]=""
        tdict[end2]=""
        if end1 in graphdict:
            tmplist = graphdict[end1]
        else:
            tmplist = []
        tmplist.append(end2)
        graphdict[end1]= tmplist

    #check and get full matrix
    for i in range(nodesize):
        if i not in tdict:
            graphdict[i]=[]

    return graphdict
# Prepare matrices # Backprop and Update # edgeList added just for benchmark testing # The goal is to make feature vectors normalized (sum equals 1), but since some feature vectors are all 0s # in those cases we'd have division by 0 so I set the min value (via np.clip) to 1. # Note: 1 is a neutral element for division i.e. it won't modify the feature vector # If there are infs that means we have a connectivity mask and 0s are where the edges in connectivity mask are, # otherwise we have an adjacency matrix and 1s symbolize the presence of edges. # active_value = 0 if np.isinf(adjacency_matrix).any() else 1 # change shape from (N,2) -> (2,N) Thresholdgraph: KNN Graph with stats one-std based methods, SingleThread version # TODO: check, only exclude large outliners # if (distMat[0,res[0][j]]<=mean+std) and (distMat[0,res[0][j]]>=mean-std): #check and get full matrix