import matplotlib.pyplot as plt
import matplotlib
from sklearn.datasets import load_iris
from sklearn.decomposition._pca import PCA
import numpy as np

class lab5:

    def __init__(self):
        pass

    def main(self):

        matplotlib.use('TkAgg')

        # Загружаем данные ирисов
        data = load_iris()
        X = data.data # признаки ирисов
        Y = data.target # сорта ирисов
        Y_str = data.target_names # названия сортов ирисов

        # setosa
        setosa_inds = Y == np.where(Y_str == "setosa")
        setosa_data = X[np.ravel(setosa_inds), :]
        # versicolor
        versicolor_inds = Y == np.where(Y_str == "versicolor")
        versicolor_data = X[np.ravel(versicolor_inds), :]
        # virginica
        virginica_inds = Y == np.where(Y_str == "virginica")
        virginica_data = X[np.ravel(virginica_inds), :]

        plt.scatter(setosa_data[:, 1], setosa_data [:, 2], marker="+", label='Setosa')
        plt.scatter(virginica_data[:, 1], virginica_data[:, 2], marker=".", label='Vergisicа')
        plt.title('Скаттерограмма Setosa и Virginica')
        plt.xlabel('Ширина чашелистика (см)')
        plt.ylabel('Ширина лепестка (см)')
        plt.legend()
        plt.show()


        # Выполняем PCA для сокращения размерности до 2D
        PCA_obj = PCA(n_components=2) # Две главные компоненты
        N = len(setosa_data)
        data_0 = np.vstack((setosa_data, virginica_data))
        reduced_data_0 = PCA_obj.fit_transform(data_0)
        print (reduced_data_0)

        #Скаттерограмма объекты сортов ирисов в пространстве главных компонент
        plt.scatter(reduced_data_0[0:N, 0], reduced_data_0[0:N, 1], marker=".", label='Setosa')
        plt.scatter(reduced_data_0[N:, 0], reduced_data_0[N:, 1], marker="*", label='Virginica')
        plt.xlabel("Главная компонента 1")
        plt.ylabel("Главная компонента 2")
        plt.title("Скатерограмма Setosa и Virginica")
        plt.legend()
        plt.show()

        # Для всех Сортов
        plt.scatter(setosa_data[:, 1], setosa_data [:, 2], marker="+", label='Setosa')
        plt.scatter(versicolor_data[:, 1], versicolor_data[:, 2], marker=".", label='Versicolor')
        plt.scatter(virginica_data[:, 1], virginica_data[:, 2], marker=".", label='Vergisicа')
        plt.title('Скаттерограмма для всех сортов ирисов')
        plt.xlabel('Ширина чашелистика (см)')
        plt.ylabel('Ширина лепестка (см)')
        plt.legend()
        plt.show()


        data = np.vstack((setosa_data, versicolor_data, virginica_data))
        reduced_data = PCA_obj.fit_transform(data)

        #Скаттерограмма объекты сортов ирисов в пространстве главных компонент
        plt.scatter(reduced_data[0:N, 0], reduced_data[0:N, 1], marker="+", label='Setosa')
        plt.scatter(reduced_data[N:2*N, 0], reduced_data[N:2*N, 1], marker=".", label='Versicolor')
        plt.scatter(reduced_data[2*N:, 0], reduced_data[2*N:, 1], marker="*", label='Virginica')
        plt.xlabel("Главная компонента 1")
        plt.ylabel("Главная компонента 2")
        plt.title("Скатерограмма всех сортов")
        plt.legend()
        plt.show()



if __name__ == "__main__":
    l5 = lab5()
    l5.main()