# Samuele Marro
# 2019-07-05

import PIL
import cv2
import numpy as np
import click
import sklearn
import sklearn.cluster
import sklearn.preprocessing
import skimage
import scipy.ndimage as ndimage
import matplotlib.pyplot as plt

@click.group()
def main():
    pass

def apply_filters(image, kernels):
    filtered_images = []
    for kernel in kernels:
        #kernel = kernel[np.newaxis] # Necessary for 3D images

        filtered = ndimage.convolve(image, kernel, mode='wrap')
        
        filtered_images.append(filtered)
    return filtered_images


def create_kernels(thetas=range(4), sigmas=(1, 3), frequencies=(0.05, 0.25)):
    kernels = []
    for theta in thetas:
        theta = theta / 4. * np.pi
        for sigma in sigmas:
            for frequency in frequencies:
                kernel = np.real(skimage.filters.gabor_kernel(frequency, theta=theta,
                                            sigma_x=sigma, sigma_y=sigma))
                                            
                kernels.append(kernel)
    return kernels


def visualise(images):
    x_size = len(images) // int(np.sqrt(len(images)).item())
    y_size = int(np.ceil(len(images) // x_size))

    print('{}x{}'.format(x_size, y_size))

    _, axarr = plt.subplots(y_size, x_size)
    for i, image in enumerate(images):
        axarr[i].imshow(image)
    plt.show()


@main.command()
@click.argument('image_path', type=click.Path(exists=True, file_okay=True, dir_okay=False))
def perform_segmentation(image_path):
    shrink = (slice(0, None, 3), slice(0, None, 3))
    final_results = []

    # Load the image
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Shrink the image
    image = image[shrink]
    final_results.append(image)

    # Convert to grayscale
    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    final_results.append(cv2.cvtColor(image, cv2.COLOR_GRAY2RGB))

    # Prepare the Gabor kernels
    kernels = create_kernels()

    # Compute the Gabor features
    filtered_images = apply_filters(image, kernels)

    # Create spatial information
    y_information, x_information = np.mgrid[0:image.shape[0], 0:image.shape[1]]

    filtered_images.append(x_information)
    filtered_images.append(y_information)

    # Create the pixel-wide features
    pixel_features = np.stack(filtered_images, axis=2)

    # Flatten the features
    flattened = pixel_features.reshape([pixel_features.shape[0] * pixel_features.shape[1], pixel_features.shape[2]])

    # Normalise to 0 mean and 1 variance
    normalised = sklearn.preprocessing.scale(flattened)

    # Apply k-means
    kmeans = sklearn.cluster.KMeans(n_clusters=8, max_iter=300)
    kmeans.fit(normalised, sample_weight=None)

    clustered = kmeans.labels_

    # Recreate the shape
    clustered_image = np.reshape(clustered, image.shape)
    final_results.append(clustered_image)

    visualise(final_results)


if __name__ == '__main__':
    main()