# This code is part of Qiskit.
#
# (C) Copyright IBM 2018, 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
iris dataset
"""

import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
from qiskit.aqua import MissingOptionalLibraryError


def iris(training_size, test_size, n, plot_data=False):
    """ returns iris dataset """
    class_labels = [r'A', r'B', r'C']
    data, target = datasets.load_iris(return_X_y=True)
    sample_train, sample_test, label_train, label_test = \
        train_test_split(data, target, test_size=1, random_state=42)

    # Now we standardize for gaussian around 0 with unit variance
    std_scale = StandardScaler().fit(sample_train)
    sample_train = std_scale.transform(sample_train)
    sample_test = std_scale.transform(sample_test)

    # Now reduce number of features to number of qubits
    pca = PCA(n_components=n).fit(sample_train)
    sample_train = pca.transform(sample_train)
    sample_test = pca.transform(sample_test)

    # Scale to the range (-1,+1)
    samples = np.append(sample_train, sample_test, axis=0)
    minmax_scale = MinMaxScaler((-1, 1)).fit(samples)
    sample_train = minmax_scale.transform(sample_train)
    sample_test = minmax_scale.transform(sample_test)

    # Pick training size number of samples from each distro
    training_input = {key: (sample_train[label_train == k, :])[:training_size]
                      for k, key in enumerate(class_labels)}
    test_input = {key: (sample_test[label_test == k, :])[:test_size]
                  for k, key in enumerate(class_labels)}

    if plot_data:
        try:
            import matplotlib.pyplot as plt
        except ImportError as ex:
            raise MissingOptionalLibraryError(
                libname='Matplotlib',
                name='iris',
                pip_install='pip install matplotlib') from ex
        for k in range(0, 3):
            plt.scatter(sample_train[label_train == k, 0][:training_size],
                        sample_train[label_train == k, 1][:training_size])

        plt.title("Iris dataset")
        plt.show()

    return sample_train, training_input, test_input, class_labels

[((28, 19, 28, 54), 'sklearn.datasets.load_iris', 'datasets.load_iris', (), '', False, 'from sklearn import datasets\n'), ((30, 8, 30, 68), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split\n'), ((43, 14, 43, 58), 'numpy.append', 'np.append', (), '', True, 'import numpy as np\n'), ((66, 8, 66, 33), 'matplotlib.pyplot.title', 'plt.title', ({(66, 18, 66, 32): '"""Iris dataset"""'}, {}), "('Iris dataset')", True, 'import matplotlib.pyplot as plt\n'), ((67, 8, 67, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((33, 16, 33, 32), 'sklearn.preprocessing.StandardScaler', 'StandardScaler', ({}, {}), '()', False, 'from sklearn.preprocessing import StandardScaler, MinMaxScaler\n'), ((38, 10, 38, 29), 'sklearn.decomposition.PCA', 'PCA', (), '', False, 'from sklearn.decomposition import PCA\n'), ((44, 19, 44, 40), 'sklearn.preprocessing.MinMaxScaler', 'MinMaxScaler', ({(44, 32, 44, 39): '(-1, 1)'}, {}), '((-1, 1))', False, 'from sklearn.preprocessing import StandardScaler, MinMaxScaler\n'), ((63, 12, 64, 74), 'matplotlib.pyplot.scatter', 'plt.scatter', ({(63, 24, 63, 73): 'sample_train[label_train == k, 0][:training_size]', (64, 24, 64, 73): 'sample_train[label_train == k, 1][:training_size]'}, {}), '(sample_train[label_train == k, 0][:training_size], sample_train\n    [label_train == k, 1][:training_size])', True, 'import matplotlib.pyplot as plt\n'), ((58, 18, 61, 53), 'qiskit.aqua.MissingOptionalLibraryError', 'MissingOptionalLibraryError', (), '', False, 'from qiskit.aqua import MissingOptionalLibraryError\n')]