#!/usr/bin/env python
# coding: utf-8

# # Graph Construction
# - We shall be creating weighted KNN graphs on pointclouds with Faiss library.
# - By default the graphs shall be directed but they can be symmetrized.
# - Our graph representation shall be same as that of torch_geometric (*i.e.* edge_index and edge_attr).
# - edge_index is a (2, E) size tensor, which contains all the edges. 
# - edge_attr is a (E,1) size tensor, which contains the scalar weights associated with the edges.
# - The graphs constructed in these notebooks *shall be used for pointcloud shape and color processing in the Primal_Dual_TV notebook and Tikhonov_reg notebook, respectively*.

# In[1]:


# import the required libs
import numpy as np
import faiss
import torch
from utilities import *
from torch_geometric.data import Data
from torch_geometric.utils import is_undirected


# In[2]:


#load the pcd/mesh 
(pos, tex, fac) = getPositionTexture("data/noisy_girl_skate.ply") # Color Processing.
#(pos, tex, fac) = getPositionTexture("data/noisy_3d_signal.ply")# Shape Processing.
pos = (pos - np.min(pos))/(np.max(pos)-np.min(pos)) #a good practice.
displaySur(**dict(position=pos, texture=tex)) # Open3d visualization works only locally!


# In[3]:


"""
- Create a knn graph using the position coords of the pointcloud.
- Assign scalar weigths to the edges.
"""

#Faiss graph construction
res = faiss.StandardGpuResources()
index = faiss.IndexFlatL2(pos.shape[1])
gpu_index_flat = faiss.index_cpu_to_gpu(res,0,index)
gpu_index_flat.add(pos.astype(np.float32))
k = 8
D, I = gpu_index_flat.search(pos.astype(np.float32),k+1)
#Convert to torch_geometric Data class 
edge_index = np.vstack((I[:,1:].flatten(), np.repeat(I[:,0].flatten(),k)))

### Shape Processing ###
#edge_attr = np.ones(edge_index.shape[1]) # Lets keep the weights equal to 1 !
### Shape Processing ###

### Color Processing ###
# RBF kernel
edge_attr = np.exp(-np.sum(((tex[I]-tex[I][:,0,None])**2), axis=2)/(0.2)**2)[:,1:].flatten() 
### Color Porcessing ###

edge_index = torch.from_numpy(edge_index).type(torch.long) # it is important to convert to torch.long
edge_attr = torch.from_numpy(edge_attr).type(torch.float32)
edge_attr = edge_attr.view(-1,1)
#getWgStats(edge_attr.numpy())


# # Graph symmetrization
# - Let $A$ be the adjacency representation of the graph. We shall be symmerizing the graph using the following transformation:
# $$ A = \frac{ A + A^{T}}{2}$$

# In[4]:


#Check if the graph is symmetric and create a temporary graph.
print(is_undirected(edge_index))
tmp_graph = Data(edge_index = edge_index, edge_attr=edge_attr, num_nodes=len(pos))

# graph symmetrization by converting to a sparse tensor.
tst = ToSparseTensor()
nG = tst(tmp_graph).adj_t.to_symmetric().to_torch_sparse_coo_tensor()
new_edge_index = torch.stack((nG.coalesce().indices()[1], nG.coalesce().indices()[0]))
new_edge_attr = nG.coalesce().values()

# Create a new graph
graph = Data(edge_index = new_edge_index.type(torch.long), edge_attr=new_edge_attr/2, 
             x=torch.from_numpy(pos).type(torch.float32), tex=torch.from_numpy(tex).type(torch.float32))
print(is_undirected(graph.edge_index))


# In[5]:


# save it
torch.save(graph, "./data/girl_skate.pt") # Color Processing
#torch.save(graph, "./data/3d_signal.pt") # Shape Processing