# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

np.random.seed(505)


# TODO: DEFINE THE CLASSIFIERS
classifiers = [
    SVC(kernel="linear"),
    SVC(kernel="rbf"),
    SVC(kernel="poly", degree=2),
    SVC(kernel="poly", degree=3),
    SVC(kernel="poly", degree=5),
]

# TODO: LOAD THE DATASET
data = np.genfromtxt('ex1_dataset.csv', delimiter=',')  # load the dataset
X = data[:, 0:-1]  # get the features
y = data [:, -1]  # get the class label
X = StandardScaler().fit_transform(X)  # perform rescaling


# Split the dataset into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=42)  # split the dataset

# CODE RELATED TO PLOTTING

# Define the names of the classifiers
names = [
    "Linear SVM",
    "RBF SVM",
    "Polynomial d2 SVM",
    "Polynomial d3 SVM",
    "Polynomial d5 SVM",
]

# for plotting
figure = plt.figure(figsize=(27, 9))
ds_cnt = 0
h = 0.02  # step size in the mesh
i = 1

x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

# just plot the dataset first
cm = plt.cm.RdBu
cm_bright = ListedColormap(["#FF0000", "#0000FF"])
ax = plt.subplot(1, len(classifiers) + 1, i)
ax.set_title("Input data")

# Plot the training points
ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_bright, edgecolors="k")

# Plot the testing points
ax.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap=cm_bright, alpha=0.6, edgecolors="k")
ax.set_xlim(xx.min(), xx.max())
ax.set_ylim(yy.min(), yy.max())
ax.set_xticks(())
ax.set_yticks(())

# ITERATE OVER CLASSIFIERS
# NOTE: clf stands for the classifier in the following code
for name, clf in zip(names, classifiers):
    ax = plt.subplot(1, len(classifiers) + 1, i)

    # TODO: FIT THE DATA AND EXTRACT SCORE
    clf.fit(X_train, y_train)  # Fit the data to the classifier
    score = clf.score(X_test, y_test)  # Extract the score from the classifier

    # Plot the decision boundary. For that, we will assign a color to each
    # point in the mesh [x_min, x_max]x[y_min, y_max].
    if hasattr(clf, "decision_function"):
        Z = clf.decision_function(np.c_[xx.ravel(), yy.ravel()])
    else:
        Z = clf.predict_proba(np.c_[xx.ravel(), yy.ravel()])[:, 1]

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    ax.contourf(xx, yy, Z, cmap=cm, alpha=0.8)

    # Plot the training points
    ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap=cm_bright, edgecolors="k")
    # Plot the testing points
    ax.scatter(
        X_test[:, 0],
        X_test[:, 1],
        c=y_test,
        cmap=cm_bright,
        edgecolors="k",
        alpha=0.6,
    )

    ax.set_xlim(xx.min(), xx.max())
    ax.set_ylim(yy.min(), yy.max())
    ax.set_xticks(())
    ax.set_yticks(())
    if ds_cnt == 0:
        ax.set_title(name)
    ax.text(
        xx.max() - 0.3,
        yy.min() + 0.3,
        ("%.2f" % score).lstrip("0"),
        size=15,
        horizontalalignment="right",
    )
    i = i + 1

plt.tight_layout()
plt.show()
