"""
This is our script to post process predictions generated by our classifier on images, as well as any other data that needs post-processing.

Further documentation found in each function.

-Blake Edwards / Dark Element
"""
import sys
import numpy as np
from scipy.sparse import csr_matrix, eye
from scipy.sparse.linalg import bicg

def normalize_adjacency_matrix(adjacency_matrix):
    """
    Helper function for denoise_predictions.

    Arguments:
        adjacency_matrix: A matrix of size h * w, 
            https://en.wikipedia.org/wiki/Adjacency_matrix
            
            For this type of normalization, should be symmetric - i.e. from an undirected graph.
            
    Returns:
        normalized_adjacency_matrix: A normalized version of our adjacency matrix, 
            such that each entry represents:
                1) If the source and dest node are connected, and if so, then
                2) how connected the source node is
                3) how connected the destination node is
        If the input is symmetric, this output will still be symmetric.
        Though it could also be normalized in other ways.
    """
    """
    We sum over all the rows of this matrix to get the degrees, 
        since it is a symmetric adjacency matrix this is the same as if we summed over the columns.

    Once we have the degrees vector, we use it as the data to create a new sparse matrix, 
        of the same shape as our adjacency matrix, and with the vector elements on the diagonal
        via using (row, col) pairs of (0,0), (1,1), ..., (h*w-1, h*w-1) with (np.arange(h*w), np.arange(h*w))
    """
    """
    We get our h*w via the length of the matrix, since we don't need individual h and w values.
    """
    adjacency_matrix_len = adjacency_matrix.shape[0]
    degrees = adjacency_matrix.sum(axis=1)
    degrees = np.array(degrees).flatten()
    degree_matrix = csr_matrix((degrees, (np.arange(adjacency_matrix_len), np.arange(adjacency_matrix_len))), shape=(adjacency_matrix_len, adjacency_matrix_len))

    """
    We use the adjacency matrix and degree matrix to make our normalized adjacency matrix, via

        M = D^(-1/2) * A * D^(-1/2)

    Where 
        D = degree matrix,
        A = adjacency matrix,
        M = normalized adjacency result

    Since our degree matrix is a diagonal matrix, 
        the elementwise matrix power of it is the same 
        as the normal matrix power, so we can quickly 
        use the built-in csr_matrix function for elementwise power
        to compute D^(-1/2). 

    We then do the dot product in the normal way.
    """
    degree_matrix = degree_matrix.power(-1./2)
    normalized_adjacency_matrix = degree_matrix.dot(adjacency_matrix.dot(degree_matrix))#DAD
    return normalized_adjacency_matrix

def denoise_predictions(predictions, neighbor_weight):
    """
    Arguments:
        predictions: np array of shape (h, w, class_n), the predictions.
            This will be referred to as an image of size h x w, as this makes 
            notation much easier.
        neighbor_weight: 
            How much importance to put on the neighbor values. 
            This could also be thought of as a smoothing factor.
            Should be between 0 and 1

    Returns:
        Each "prediction" is an element in our predictions input,
            and is a vector of length class_n. There are h * w of these predictions
            in our entire predictions array.

        For each prediction vector in our predictions, 
            we produce a new denoised prediction vector,
            taking into account both the nearby prediction vectors to this element,
            and the original prediction vector.

        Returns a new predictions array, of the same shape (h, w, class_n) as the original.
    """
    """
    First we need to create an adjacency matrix for each entry in our predictions array, 
        as if each node was a prediction vector, so that the graph is a rectangle of shape h x w,
        and each node is connected to all of it's nearby nodes. 

    Ex, with h = 3, w = 4:

    . . . .
    . . . . (I think it may be easier to just imagine the lines, due to limitations of text art)
    . . . .
    
    Such that the minimum amount of connections a node will have (unless there is a really really small h and w)
        will be 3, and the maximum will be 8.

    Since adjacency matrices are sparse (mostly zeros), we can save memory by using different compression algorithms to 
        only store the nonzero elements. I'll be using scipy's CSR element, as this algorithm offers easy matrix-vector 
        dot products, and as a result of this can be extended to matrix-matrix dot products - CSR is efficient when
        you need to perform linear algebra ops on your sparse matrices.

    Since the adjacency matrix will be of size (h*w, h*w), it quickly becomes absolutely necessary to use a compression format,
        or else numpy won't even be able to hold it in memory. For an example of h = 300, w = 400, the raw adjacency matrix
        would contain h*w * h*w elements, = 300*400 * 300*400 = 120000 * 120000 = 1.44e10 .

    Due to the simple nature of our graph, undirected and where each node is connected in a simple and deterministic manner to its neighbors,
        we can construct it with O(n) complexity. I designed a way to incrementally create each row in this matrix given the h and w,
        as this is all that's necessary given our knowledge of how the graph is constructed. I could loop through every element and construct it
        the way it would be done by hand, however this would have O(n^2) complexity, and since n = h*w, that is really really bad. For the sake
        of optimizing down to O(n), I had to observe and replicate the patterns in adjacency matrices of this type. It is not as obviously intuitive 
        as the by-hand method, however it will run at O(n), able to loop through and construct rows of the adjacency matrix at a time.
    """
    """
    Prepare these lists for constructing our CSR matrix
    """
    rows = []
    cols = []
    data = []

    """
    Get this data from our shape
    """
    h, w, class_n = predictions.shape

    def ins(i, j, x):
        """
        Arguments:
            i: row index of x
            j: col index of x
            x: data to add
        Returns:
            Adds i, our row index, to the row indices list
            Adds j, our col index, to the col indices list
            Adds x, our data, to the data list

            This is just a quick function to save time in the next steps.
        """
        if j >= 0 and j < h*w:
            rows.append(i)
            cols.append(j)
            data.append(x)


    """
    Loop through and construct our rows
    """
    for i in range(h*w):
        if i % w == w-1:
            ins(i, i-(w+1), 1)
            ins(i, i-(w), 1)
            ins(i, i-1, 1)
            ins(i, i+(w-1), 1)
            ins(i, i+(w), 1)

        elif i % w == 0:
            ins(i, i-(w), 1)
            ins(i, i-(w-1), 1)
            ins(i, i+1, 1)
            ins(i, i+(w), 1)
            ins(i, i+(w+1), 1)

        else:
            ins(i, i-(w+1), 1)
            ins(i, i-(w), 1)
            ins(i, i-(w-1), 1)
            ins(i, i-1, 1)
            ins(i, i+1, 1)
            ins(i, i+(w-1), 1)
            ins(i, i+(w), 1)
            ins(i, i+(w+1), 1)

    """
    Now that lists are finished changing in size, 
        we cast to nparrays
    """
    rows = np.array(rows)
    cols = np.array(cols)
    data = np.array(data)

    """
    Using our data and row and column indices, we create our adjacency matrix as a CSR Sparse matrix
    """
    adjacency_matrix = csr_matrix((data, (rows, cols)), shape=(h*w, h*w))

    """
    We then normalize this adjacency matrix
    """
    normalized_adjacency_matrix = normalize_adjacency_matrix(adjacency_matrix)

    """
    Then we use the following equation, from this paper:

        http://journals.plos.org/plosone/article?id=10.1371/journal.pone.0051947

        (I - l*M)*f = (1-l)*y

    In order to solve for new values of y (our predictions), so that we have
        new predictions in f.

    That paper has the explanation for that equation. 
    Once we have that equation, we can simplify so that:

        A = (I - l * M)
        x = f
        b = (1-l)*y

    Our equation is now:

        Ax = b ,

    Which is just the setup for a least squares approximation of x:

        https://en.wikipedia.org/wiki/Least_squares
        https://www.youtube.com/watch?v=MC7l96tW8V8 (Khan Academy)

    So we prepare those components first.
    """
    A = eye(h*w) - neighbor_weight * normalized_adjacency_matrix

    predictions = np.reshape(predictions, (-1, class_n))
    b = (1.-neighbor_weight) * predictions

    """
    Unfortunately, least squares usually expects x and b to be vectors.
    Right now, x is going to be of shape (h*w, class_n), and b is also of shape (h*w, class_n).
    These are very much not vectors, but matrices instead. So how can we extend least squares to work in this setting?
    
    It's actually very simple, due to one property.

    For the system Ax = b, where 
        A is a matrix of size n x n
        x is a matrix of size n x m
        b is a matrix of size n x m,

    Let us look at a column vector in matrix b, at position i. 
    Each element in this vector is determined entirely by a column vector in matrix x,
        at position i,
        as well as the entire matrix A.
    For an understanding of this, think about how the elements of this column vector in b are
        obtained through the dot product of A and x.
    
    What this means is that we can solve for the entire matrix x by solving for each of the columns
        in x independently, looping across both the columns we are solving for in x and the other matching column
        in b, and solving using normal least squares.

    Remember that x is our new prediction values, so this is the last math-heavy step.
    """
    x = np.zeros(b.shape, dtype=np.float32)
    for i in range(class_n):
        col = b[:,i]
        col = np.reshape(col, (-1, 1))
        """
        Solve for this column using Bi-Conjugate Gradient Iteration from scipy,
            and set the result column into our resulting x matrix.
        """
        x[:,i], flags = bicg(A, col)
    
    """
    At this point our entire x result matrix has been constructed, 
        however it is of shape (h*w, class_n).
    So we reshape it into (h, w, class_n), and rename it to the more
        appropriate denoised_predictions - since these are our denoised predictions.
    """
    denoised_predictions = np.reshape(x, (h, w, class_n))
    return denoised_predictions
    

