import numpy as np

from matplotlib import pyplot as plt


def display_image(img_matrix: np.matrix, size: int = 1):
    """
    Takes a matrix of values and displays it as a grayscale image.

    Args:
        img_matrix (np.matrix):     The matrix that is to be displayed as image
        size (int):                 The size of the window
    """
    plt.figure(figsize=(size, size))
    plt.imshow(img_matrix, interpolation='nearest', cmap='gray')
    plt.show()
    plt.pause(0.1)


def visualize_image(X: np.matrix, y: np.matrix = None, size: int = 20, transpose: bool = True, predictions=None, feed_values=None):
    """
    Takes a set and displays each test case as image (20x20) and displays the neural network's prediction.

    Args:
        X (np.matrix):      The set that is to be visualized
        y (np.matrix):      (optional) The set's corresponding output.
        size (int):         (optional) Size of the image (original image)
        transpose (bool):   (optional) Whether or not to transpose the image
        predictions:        (optional) The set's prediction made by the neural network.
        feed_values:        (optional) The feed forward values generated by the neural network when passing X as input.
    """
    idx = np.random.permutation(len(X))

    i = 0
    plt.ion()
    input_key = ''
    while not input_key == 'q' and i < len(X):
        img = np.array(X[idx[i]]).reshape((size, size))

        if transpose:
            img = img.T

        display_image(img)

        if predictions is not None:
            print('\n{:<14s}'.format('Prediction:'), '{:^2d}'.format((int(np.where(predictions[idx[i]] == 1)[0]) + 1) % 10))
        if y is not None:
            print('{:<14s}'.format('Value:'), '{:^2d}'.format((int(np.where(y[idx[i]] == 1)[0]) + 1) % 10))
        if feed_values is not None:
            print('Confidence:', feed_values[len(feed_values)-1][idx[i]][(int(np.where(predictions[idx[i]] == 1)[0]))])

        input_key = input('\rPress Enter to continue or \'q\' to exit....')
        plt.close()
        i += 1

    plt.ioff()


def visualize_training_step(iteration, **kwargs):
    """
    Visualizes what an iteration (or possibly multiple iterations if step > 1) of the gradient descent algorithm do
    in terms of improving the fit of the hypothesis to the data.
    
    Args:
        iteration:   The current iteration of the gradient descent algorithm
        kwargs:      Additional settings: 'step' -> after how many steps the plot should be updated
    """
    step = kwargs['step'] if 'step' in kwargs else 1
    theta = kwargs['theta'] if 'theta' in kwargs else None
    X = kwargs['X'] if 'X' in kwargs else None

    if theta is None or X is None:
        raise Exception('Theta or X were not passed on to the callback function! Check your gd_params.callback_args')

    if iteration % step == step-1:
        plt.plot(np.ravel(X[0:, 0].T), X.dot(theta.T), color='red', linewidth=0.5)
        plt.pause(0.2)


def visualize_final_result(X: np.matrix, y: np.matrix, theta=None, pause: float = 1.0):
    """
    Takes a set X and a corresponding y and scatters the data. The first column of the set X acts as the x-axis values.
    Additionally a parameter vector theta can be passed on so that the hypothesis is also plotted.

    Args:
        X (np.matrix):      A set X that is to be scattered (Shape: MxN)
        y (np.matrix):      A set y corresponding to the output of X (Shape: Mx1)
        theta:              (optional) A parameter vector (Shape: 1xN)
        pause (float):      (optional) Determines how long should be waited until the hypothesis is plotted 
    """
    plt.ion()
    plt.scatter(np.ravel(X[0:, 0].T), np.ravel(y.T))
    plt.show()
    if theta is not None:
        plt.pause(pause)
        plt.plot(np.ravel(X[0:, 0].T), X.dot(theta.T), color='red')

    input_key = 'no_key'
    while input_key == 'no_key':
        plt.pause(0.05)
        input_key = input('\rPress any key to continue...')