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

# # Adjacency Matrix

import numpy as np

from scipy.sparse import csr_matrix, lil_matrix, coo_matrix


class adj_matrix(object):
    def __init__(self, symmetric_indices=False):
        
        self._sym = symmetric_indices
        if not symmetric_indices:
            self._x_labels = []
            self._y_labels = []

            self._x_label_index={}
            self._y_label_index={}
        
        else:
            self._labels = []
            self._label_index={}
        
        self._x = []
        self._y = []
        self._data = []
        
        self._mat = None
    
    def _get_ix(self, label):
        i = self._x_label_index.get(label)
        if i is None:
            i = len(self._x_labels)
            self._x_labels.append(label)
            self._x_label_index[label] = i
        return i
    
    def _get_iy(self, label):
        i = self._y_label_index.get(label)
        if i is None:
            i = len(self._y_labels)
            self._y_labels.append(label)
            self._y_label_index[label] = i
        return i
    
    def _get_i(self, label):
        i = self._label_index.get(label)
        if i is None:
            i = len(self._labels)
            self._labels.append(label)
            self._label_index[label] = i
        return i
    
    def add_entry(self, x, y, data):
        
        if self._sym:
            ix = self._get_i(x)
            iy = self._get_i(y)
        
        else:
            ix = self._get_ix(x)
            iy = self._get_iy(y)
        
        self._x.append(ix)
        self._y.append(iy)
        self._data.append(data)
    
    def compile_to_mat(self):
        if self._sym:
            sx = len(self._labels)
            sy = len(self._labels)
        else:
            sx = len(self._x_labels)
            sy = len(self._y_labels)
        
        self._mat = coo_matrix((self._data, (self._x, self._y)), shape=(sx,sy))
        return self._mat
    
    def get_csr(self):
        return self.compile_to_mat().tocsr()
    
    def get_labels(self):
        if self._sym:
            return self._labels
        return self._x_labels, self._y_labels




