"""
This file contains some utility functions to use in conjunction with SFA.
"""
import numpy as np

def randomWalkFromCSC(csc_matrix,
                      n_steps=10, starting_node=None, restart_rate=0):
    """
    This method produces an index trajectory on a graph by random walk,
    when the adjacency matrix is given in Compressed Sparse Column (CSC)
    Matrix format as implemented by scipy.
    The trajectory can be used to emulate spectral embeddings with SFA.

    Parameters
    ----------
    osc_matrix : sparse matrix (csc) of shape (n_nodes, n_nodes)
        This CSC matrix should contain the adjacency matrix of a graph.

    n_steps : int, default=10
        The number of overall steps of the random walk.

    starting_node : int or None, (default None)
        The index of the node from where to start the random walk.
        If None :
            A random node will be sampled between 0 and the maximum
            node index.

    restart_rate : int, default=0
        After which step will the random walk be restarted from a
        random node. This can be used to get better coverage of a
        graph with large diameter.
        If 0 :
            The trajectory will be generated by one consecutive random
            walk.
    """
    node_trajectory = []
    if starting_node is None:
        current_node = np.random.randint(0, csc_matrix.shape[0])
    else:
        current_node = starting_node
    for step_idx in range(n_steps):
        node_trajectory.append(current_node)
        index_beginning = csc_matrix.indptr[current_node]
        index_end = csc_matrix.indptr[current_node+1]
        neighbor_indices = csc_matrix.indices[index_beginning:index_end]
        current_node = np.random.choice(neighbor_indices)
        if restart_rate != 0:
            if step_idx % restart_rate == 0:
                current_node = np.random.randint(0, csc_matrix.shape[0])
    return node_trajectory
