import numpy as np
import scipy
from scipy.linalg import expm
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
import pandas as pd
from qiskit import BasicAer
from qiskit.aqua.input import ClassificationInput
from qiskit.aqua import run_algorithm


def Currency(training_size, test_size, n, PLOT_DATA):
    class_labels = [r'Buy', r'Sell', r'Hold']
    training_url = "https://drive.google.com/file/d/1dGr9w629hZMlHcTvNw1sGlhm3J57osze/view?usp=sharing"
    training=pd.read_csv(training_url, sep=',',header=0)
    trainingNP = training.to_numpy()
    sample_train = trainingNP[:,:-2]
    label_train = trainingNP[:,-2]
    label_train[label_train == 'Buy'] = 0
    label_train[label_train == 'Sell'] = 1
    label_train[label_train == 'Hold'] = 2
    
    test_url = "https://drive.google.com/file/d/1wsyURMe0wRXysJvLn1T9bXIsMqd8MocG/view?usp=sharing"
    test=pd.read_csv(test_url, sep=',',header=0)
    testNP = test.to_numpy()
    sample_test = testNP[:,:-2]
    label_test = testNP[:,-2]
    label_test[label_test == 'Buy'] = 0
    label_test[label_test == 'Sell'] = 1
    label_test[label_test == 'Hold'] = 2
    
    # Now we standarize 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:
        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("PCA dim. reduced Currency dataset")
        plt.show()

    return sample_train, training_input, test_input, class_labels

n = 2  # dimension of each data point
sample_Total, training_input, test_input, class_labels = Currency(
    training_size=80,
    test_size=20, n=n, PLOT_DATA=True
)
temp = [test_input[k] for k in test_input]
total_array = np.concatenate(temp)

aqua_dict = {
    'problem': {'name': 'classification', 'random_seed': 10598},
    'algorithm': {
        'name': 'QSVM'
    },
    'feature_map': {'name': 'SecondOrderExpansion', 'depth': 2, 'entangler_map': [[0, 1]]},
    'multiclass_extension': {'name': 'AllPairs'},
    'backend': {'shots': 1024}
}

backend = BasicAer.get_backend('qasm_simulator')
algo_input = ClassificationInput(training_input, test_input, total_array)
result = run_algorithm(aqua_dict, algo_input, backend=backend)
for k,v in result.items():
    print("'{}' : {}".format(k, v))