# 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, 'sklearn.datasets.load_iris', 'datasets.load_iris', False, 'from sklearn import datasets\n'), (30, 'sklearn.model_selection.train_test_split', 'train_test_split', False, 'from sklearn.model_selection import train_test_split\n'), (43, 'numpy.append', 'np.append', True, 'import numpy as np\n'), (66, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (67, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (33, 'sklearn.preprocessing.StandardScaler', 'StandardScaler', False, 'from sklearn.preprocessing import StandardScaler, MinMaxScaler\n'), (38, 'sklearn.decomposition.PCA', 'PCA', False, 'from sklearn.decomposition import PCA\n'), (44, 'sklearn.preprocessing.MinMaxScaler', 'MinMaxScaler', False, 'from sklearn.preprocessing import StandardScaler, MinMaxScaler\n'), (63, 'matplotlib.pyplot.scatter', 'plt.scatter', True, 'import matplotlib.pyplot as plt\n'), (58, 'qiskit.aqua.MissingOptionalLibraryError', 'MissingOptionalLibraryError', False, 'from qiskit.aqua import MissingOptionalLibraryError\n')]