import os
import numpy as np
import scipy.sparse as sp
from sklearn.preprocessing import MultiLabelBinarizer, LabelBinarizer, normalize

def load_npz_to_sparse_graph(file_name):
    """Load a SparseGraph from a Numpy binary file.

    Parameters
    ----------
    file_name : str
        Name of the file to load.

    Returns
    -------
    sparse_graph : SparseGraph
        Graph in sparse matrix format.

    """
    with np.load(file_name) as loader:
        loader = dict(loader)
        adj_matrix = sp.csr_matrix((loader['adj_data'], loader['adj_indices'], loader['adj_indptr']),
                                   shape=loader['adj_shape'])

        if 'attr_data' in loader:
            # Attributes are stored as a sparse CSR matrix
            attr_matrix = sp.csr_matrix((loader['attr_data'], loader['attr_indices'], loader['attr_indptr']),
                                        shape=loader['attr_shape'])
        elif 'attr_matrix' in loader:
            # Attributes are stored as a (dense) np.ndarray
            attr_matrix = loader['attr_matrix']
        else:
            attr_matrix = None

        if 'labels_data' in loader:
            # Labels are stored as a CSR matrix
            labels = sp.csr_matrix((loader['labels_data'], loader['labels_indices'], loader['labels_indptr']),
                                   shape=loader['labels_shape'])
        elif 'labels' in loader:
            # Labels are stored as a numpy array
            labels = loader['labels']
        else:
            labels = None

        node_names = loader.get('node_names')
        attr_names = loader.get('attr_names')
        class_names = loader.get('class_names')
        metadata = loader.get('metadata')

    return SparseGraph(adj_matrix, attr_matrix, labels, node_names, attr_names, class_names, metadata)

def SparseGraph(adj_matrix, attr_matrix, labels, node_names, attr_names, class_names, metadata):
    if sp.isspmatrix(adj_matrix):
        adj_matrix = adj_matrix.tocsr().astype(np.float32)
    else:
        raise ValueError("Adjacency matrix must be in sparse format (got {0} instead)"
                         .format(type(adj_matrix)))

    if adj_matrix.shape[0] != adj_matrix.shape[1]:
        raise ValueError("Dimensions of the adjacency matrix don't agree")

    if attr_matrix is not None:
        if sp.isspmatrix(attr_matrix):
            attr_matrix = attr_matrix.tocsr().astype(np.float32)
        elif isinstance(attr_matrix, np.ndarray):
            attr_matrix = attr_matrix.astype(np.float32)
        else:
            raise ValueError("Attribute matrix must be a sp.spmatrix or a np.ndarray (got {0} instead)"
                             .format(type(attr_matrix)))

        if attr_matrix.shape[0] != adj_matrix.shape[0]:
            raise ValueError("Dimensions of the adjacency and attribute matrices don't agree")

    if labels is not None:
        if labels.shape[0] != adj_matrix.shape[0]:
            raise ValueError("Dimensions of the adjacency matrix and the label vector don't agree")

    if node_names is not None:
        if len(node_names) != adj_matrix.shape[0]:
            raise ValueError("Dimensions of the adjacency matrix and the node names don't agree")

    if attr_names is not None:
        if len(attr_names) != attr_matrix.shape[1]:
            raise ValueError("Dimensions of the attribute matrix and the attribute names don't agree")

    return get_dataset(adj_matrix, attr_matrix, labels)

def get_dataset(adj_matrix, attr_matrix, labels):
    labels = binarize_labels(labels)
    # convert to binary bag-of-words feature representation if necessary
    if not is_binary_bag_of_words(attr_matrix):
        attr_matrix = to_binary_bag_of_words(attr_matrix)

    # some assertions that need to hold for all datasets
    # adj matrix needs to be symmetric
    # assert (adj_matrix != adj_matrix.T).nnz == 0
    # features need to be binary bag-of-word vectors
    # assert is_binary_bag_of_words(attr_matrix), f"Non-binary node_features entry!"

    return adj_matrix, attr_matrix, labels


def binarize_labels(labels, sparse_output=False, return_classes=False):
    """Convert labels vector to a binary label matrix.

    In the default single-label case, labels look like
    labels = [y1, y2, y3, ...].
    Also supports the multi-label format.
    In this case, labels should look something like
    labels = [[y11, y12], [y21, y22, y23], [y31], ...].

    Parameters
    ----------
    labels : array-like, shape [num_samples]
        Array of node labels in categorical single- or multi-label format.
    sparse_output : bool, default False
        Whether return the label_matrix in CSR format.
    return_classes : bool, default False
        Whether return the classes corresponding to the columns of the label matrix.

    Returns
    -------
    label_matrix : np.ndarray or sp.csr_matrix, shape [num_samples, num_classes]
        Binary matrix of class labels.
        num_classes = number of unique values in "labels" array.
        label_matrix[i, k] = 1 <=> node i belongs to class k.
    classes : np.array, shape [num_classes], optional
        Classes that correspond to each column of the label_matrix.

    """
    if hasattr(labels[0], '__iter__'):  # labels[0] is iterable <=> multilabel format
        binarizer = MultiLabelBinarizer(sparse_output=sparse_output)
    else:
        binarizer = LabelBinarizer(sparse_output=sparse_output)
    label_matrix = binarizer.fit_transform(labels).astype(np.float32)
    return (label_matrix, binarizer.classes_) if return_classes else label_matrix

def get_train_val_test_split(random_state,
                             labels,
                             train_examples_per_class=None, val_examples_per_class=None,
                             test_examples_per_class=None,
                             train_size=None, val_size=None, test_size=None):
    num_samples, num_classes = labels.shape
    remaining_indices = list(range(num_samples))

    if train_examples_per_class is not None:
        train_indices = sample_per_class(random_state, labels, train_examples_per_class)
    else:
        # select train examples with no respect to class distribution
        train_indices = random_state.choice(remaining_indices, train_size, replace=False)
        train_mask = sample_mask(train_indices, labels.shape[0])

    if val_examples_per_class is not None:
        val_indices = sample_per_class(random_state, labels, val_examples_per_class, forbidden_indices=train_indices)
    else:
        remaining_indices = np.setdiff1d(remaining_indices, train_indices)
        val_indices = random_state.choice(remaining_indices, val_size, replace=False)
        val_mask = sample_mask(val_indices, labels.shape[0])

    forbidden_indices = np.concatenate((train_indices, val_indices))
    if test_examples_per_class is not None:
        test_indices = sample_per_class(random_state, labels, test_examples_per_class,
                                        forbidden_indices=forbidden_indices)
    elif test_size is not None:
        remaining_indices = np.setdiff1d(remaining_indices, forbidden_indices)
        test_indices = random_state.choice(remaining_indices, test_size, replace=False)
        test_mask = sample_mask(test_indices, labels.shape[0])
    else:
        test_indices = np.setdiff1d(remaining_indices, forbidden_indices)
        test_mask = sample_mask(test_indices, labels.shape[0])

    # assert that there are no duplicates in sets
    assert len(set(train_indices)) == len(train_indices)
    assert len(set(val_indices)) == len(val_indices)
    assert len(set(test_indices)) == len(test_indices)
    # assert sets are mutually exclusive
    assert len(set(train_indices) - set(val_indices)) == len(set(train_indices))
    assert len(set(train_indices) - set(test_indices)) == len(set(train_indices))
    assert len(set(val_indices) - set(test_indices)) == len(set(val_indices))
    if test_size is None and test_examples_per_class is None:
        # all indices must be part of the split
        assert len(np.concatenate((train_indices, val_indices, test_indices))) == num_samples

    if train_examples_per_class is not None:
        train_labels = labels[train_indices, :]
        train_sum = np.sum(train_labels, axis=0)
        # assert all classes have equal cardinality
        assert np.unique(train_sum).size == 1

    if val_examples_per_class is not None:
        val_labels = labels[val_indices, :]
        val_sum = np.sum(val_labels, axis=0)
        # assert all classes have equal cardinality
        assert np.unique(val_sum).size == 1

    if test_examples_per_class is not None:
        test_labels = labels[test_indices, :]
        test_sum = np.sum(test_labels, axis=0)
        # assert all classes have equal cardinality
        assert np.unique(test_sum).size == 1

    y_train = np.zeros(labels.shape)
    y_val = np.zeros(labels.shape)
    y_test = np.zeros(labels.shape)
    y_train[train_mask, :] = labels[train_mask, :]
    y_val[val_mask, :] = labels[val_mask, :]
    y_test[test_mask, :] = labels[test_mask, :]

    return y_train, y_val, y_test, train_mask, val_mask, test_mask

def sample_per_class(random_state, labels, num_examples_per_class, forbidden_indices=None):
    num_samples, num_classes = labels.shape
    sample_indices_per_class = {index: [] for index in range(num_classes)}

    # get indices sorted by class
    for class_index in range(num_classes):
        for sample_index in range(num_samples):
            if labels[sample_index, class_index] > 0.0:
                if forbidden_indices is None or sample_index not in forbidden_indices:
                    sample_indices_per_class[class_index].append(sample_index)

    # get specified number of indices for each class
    return np.concatenate(
        [random_state.choice(sample_indices_per_class[class_index], num_examples_per_class, replace=False)
         for class_index in range(len(sample_indices_per_class))
         ])

def is_binary_bag_of_words(features):
    features_coo = features.tocoo()
    return all(single_entry == 1.0 for _, _, single_entry in zip(features_coo.row, features_coo.col, features_coo.data))

def to_binary_bag_of_words(features):
    """Converts TF/IDF features to binary bag-of-words features."""
    features_copy = features.tocsr()
    features_copy.data[:] = 1.0
    return features_copy

def sample_mask(idx, l):
    """Create mask."""
    mask = np.zeros(l)
    mask[idx] = 1
    return np.array(mask, dtype=np.bool)

def load_npz_data(dataset_str, seed):
    adj, features, labels = load_npz_to_sparse_graph("data/npz/{}.npz".format(dataset_str))
    random_state = np.random.RandomState(seed)
    y_train, y_val, y_test, train_mask, val_mask, test_mask = get_train_val_test_split(random_state, labels, train_size=20*labels.shape[1], val_size=30*labels.shape[1])
    return adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask
