import numpy as np
import matplotlib.pyplot as plt
import time, copy, warnings


from sklearn import linear_model, svm
from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
from sklearn.neural_network import MLPClassifier

from sklearn.model_selection import train_test_split

from sklearn.metrics import classification_report, confusion_matrix, accuracy_score


size = 4                            # size of image

binaries = np.array([[int(i) for i in np.binary_repr(a, size)] for a in range(1, 2**size - 1)]) 


#--------------------------------------------------------------------------------------------------

ones = np.ones((size, size))

X = []                              # X collects (size x size) feature matrices
y = []                              # y collects class labels 1 or 0

for b in binaries:
    x = b*ones
    
    X.append(x.flatten())           # image with vertical strips, class 1
    y.append(1)
    
    X.append((x.T).flatten())       # image with horizontal strips, class 0  
    y.append(0)

X, y = np.array(X), np.array(y)

#--------------------------------------------------------------------------------------------------    


num_examples, num_features = X.shape
print(f'number of examples = {num_examples}')    
print(f'number of features = {num_features}')    

num_features == size*size

idx = np.random.choice(num_examples)

x = X[idx].reshape(size,size)
print(f'x_{idx} of size {x.shape} :')
plt.imshow(x, cmap='gray', vmin=0, vmax=1)
plt.show()

if y[idx]==1:
    class_ = 'vertical'
elif y[idx]==0:
    class_ = 'horizontal'

print(f'true label = y = {y[idx]} = image of {class_} strips \n')


X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=0)


num_train_examples, num_test_examples = X_train.shape[0], X_test.shape[0]
print(f'number of train examples = {num_train_examples}')    
print(f'number of test examples = {num_test_examples}\n') 


print(f'number of features = {num_features}')

num_train_examples + num_test_examples == num_examples

from qiskit.circuit.library import PauliFeatureMap, ZFeatureMap, ZZFeatureMap

'For angle encoding, num_qubits = num_features'


#fm = PauliFeatureMap(feature_dimension=num_features, reps=1, paulis=['Z']) 
fm = ZZFeatureMap(feature_dimension=num_features, reps=1) 

fm.decompose().draw() 


from qiskit_machine_learning.circuit.library import RawFeatureVector

'For amplitude encoding, num_qubits = n = np.log2(num_features)''

fm = RawFeatureVector(feature_dimension=n) 
fm.draw() 


from qiskit.circuit.library import RealAmplitudes


'''
For angle encoding, num_qubits = num_features
For amplitude encoding, num_qubits = n = np.log2(num_features)
'''

pqc = RealAmplitudes(num_qubits=num_features, reps=3)
pqc.decompose().draw()


from qiskit.algorithms.optimizers import COBYLA

# This defines where to train our classifier: on a simulator or a real quantum computer
from qiskit.primitives import Sampler                             # here we have chosen a simulator 

from qiskit_machine_learning.algorithms.classifiers import VQC    # !pip install qiskit[machine-learning]


#------------------------------------------------------------------------------------------------

optimizer = COBYLA(maxiter=60)                 # maxiter = Maximum number of iterations (function evaluations)
sampler = Sampler()

#------------------------------------------------------------------------------------------------

from IPython.display import clear_output
objective_func_vals = []

                                    # a callback that can access the intermediate data during the optimization
def callback(weights, obj_func_eval):
    clear_output(wait=True)
    objective_func_vals.append(obj_func_eval)
    plt.xlabel("Iteration")
    plt.ylabel("Objective function value")
    plt.plot(range(len(objective_func_vals)), objective_func_vals)
    plt.show()

#------------------------------------------------------------------------------------------------

vqc = VQC(
    sampler=sampler,
    feature_map=fm,
    ansatz=pqc,
    optimizer=optimizer,
    callback=callback,
)

%%time

'''Training time seems to increase with the number of features, examples, and learnable-parameters'''

vqc.fit(X_train, y_train)

%%time

def performance_VQC(X, y, name):
    y_pred = vqc.predict(X)
    print(classification_report(y, y_pred))
    print(f'{confusion_matrix(y, y_pred)} = confusion matrix for {name} set \n')
    return y_pred
    

y_pred_train = performance_VQC(X_train, y_train, 'train')
y_pred_test = performance_VQC(X_test, y_test, 'test')


idx = np.random.choice(len(X_test))

x = X_test[idx]
print(f'x of size {(size,size)} :')
plt.imshow(x.reshape(size,size), cmap='gray')
plt.show()

if y[idx]==1:
    class_ = 'vertical'
elif y[idx]==0:
    class_ = 'horizontal'

print(f'true label = y = {y[idx]} = image with {class_} strips \n')


y_pred_ = vqc.predict(x)
print(f'predicted label = {y_pred_}\n')


%%time
warnings.filterwarnings('ignore')

def train_test(ML_model):
    print(20*' ', ML_model,'\n')
    ML_model.fit(X_train, y_train)                         # fit model on train set
    
    y_train_pred = ML_model.predict(X_train)               # predict with the trained model
    y_test_pred = ML_model.predict(X_test)
    
    
    print(classification_report(y_train, y_train))         # evaluate model performance 
    print(confusion_matrix(y_train, y_train_pred), "= confusion matrix for train set \n")

    print(classification_report(y_test, y_test_pred))
    print(confusion_matrix(y_test, y_test_pred), "= confusion matrix for test set \n")
    
    print(70*'=')
    

#----------------------------------------------------------------------------------

print(70*'=')
ML_model = linear_model.LogisticRegression()                # create model instance
train_test(ML_model)


ML_model = MLPClassifier()                            
train_test(ML_model)


ML_model = svm.SVC() 
train_test(ML_model)


ML_model = RandomForestClassifier()
train_test(ML_model)
