"""
Script used to visualize results from any of the techniques

authors : Rohit Bhattacharya, Azwad Sabik
emails  : rohit.bhattachar@gmail.com, azwadsabik@gmail.com
"""

# imports
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def visualize2d(X, Y, output_name):
    '''
    Visualizes in 2d
    '''
    
    # make a colour map that colours the points
    # based on their unique labels
    unique_labels = set(Y)
    num_unique_labels = len(unique_labels)
    unique_labels_dict = dict(zip(unique_labels, range(num_unique_labels)))
    colour_map = [unique_labels_dict[l] for l in Y]

    # save and plot
    plt.scatter(X[:,0], X[:,1], c=colour_map)
    plt.savefig(output_name)
    plt.show()

def visualize3d(X, Y, output_name):
    '''
    Visualizes in 3d
    '''

    # make a colour map that colours the points
    # based on their unique labels
    unique_labels = set(Y)
    num_unique_labels = len(unique_labels)
    unique_labels_dict = dict(zip(unique_labels, range(num_unique_labels)))
    colour_map = [unique_labels_dict[l] for l in Y]

    # save and plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(X[:,0], X[:,1], X[:,2], c=colour_map)
    plt.savefig(output_name)
    plt.show()
