#!/usr/bin/env python
# -*- coding: utf-8 -*-


import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
import spectral_embedding

import sys
sys.path.insert(0, '../cluster')
import kmeans


def test():
    k = 2
    X, y_true = make_moons(n_samples=500, random_state=0, noise=0.01)
    Y = spectral_embedding.transform(X, k, n_neighbors=7, sigma=0.1)
    n = np.linalg.norm(Y, axis=1)
    n = n.reshape(-1, 1)
    Y = Y / n
    # Apply K-Means to cluster Y
    y_pred, _, _ = kmeans.kmeans(Y, k)

    fig = plt.figure()
    ax = fig.add_subplot(121)
    ax.scatter(np.arange(len(Y)), Y[:, 0])
    ax.set_title("Eigenvector 1")
    ax = fig.add_subplot(122)
    ax.scatter(np.arange(len(Y)), Y[:, 1])
    ax.set_title("Eigenvector 2")

    # Plot the data
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(X[y_true==0, 0], X[y_true==0, 1], c='b', alpha=0.5, label="Class 1")
    ax.scatter(X[y_true==1, 0], X[y_true==1, 1], c='g', alpha=0.5, label="Class 2")
    ax.set_title("Original data")
    ax.legend()

    # Plot the predictions
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(X[y_pred==0, 0], X[y_pred==0, 1], c='r', alpha=0.5, label="Class 1")
    ax.scatter(X[y_pred==1, 0], X[y_pred==1, 1], c='y', alpha=0.5, label="Class 2")
    ax.set_title("Result of clustering")
    ax.legend()

    # Plot the transformed data
    fig = plt.figure()
    ax = fig.add_subplot(111)
    idx_class0 = np.argwhere(y_true==0)
    idx_class1 = np.argwhere(y_true==1)
    ax.scatter(Y[idx_class0, 0], Y[idx_class0, 1], c='b', alpha=0.5, label="Class 1")
    ax.scatter(Y[idx_class1, 0], Y[idx_class1, 1], c='g', alpha=0.5, label="Class 2")
    ax.set_title("Original data after spectral embedding")
    ax.legend()

    print("Number in class 0: {}".format(np.sum(y_pred==0)))
    print("Number in class 1: {}".format(np.sum(y_pred==1)))

    plt.show()


if __name__ == '__main__':
    test()
