# pylint: disable=import-error, wrong-import-position, too-few-public-methods, no-self-use, function-redefined
from functools import partial
from typing import Sequence
import os
import sys

import numpy as np
import matplotlib.pyplot as plt

# Import Qiskit packages
from qiskit.circuit.library import ZZFeatureMap
from qiskit_machine_learning.kernels import TrainableFidelityQuantumKernel
from qiskit_machine_learning.utils.loss_functions import KernelLoss

feature_map = ZZFeatureMap(3)
user_params = [feature_map.parameters[1]]

qkernel = TrainableFidelityQuantumKernel(
    feature_map=feature_map,
    training_parameters=user_params,
)

X_train = [[0.6, 0.2], [0.5, 0.3], [0.3, 0.7], [0.1, 0.5]]
y_train = np.array([0, 0, 1, 1])

class CustomKernelLoss(KernelLoss):
    """Example Kernel Loss class"""

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    # Evaluate the Loss of a trainable quantum kernel
    # at a particular setting of user_param_values on
    # a particular dataset.
    def evaluate(
        self,
        parameter_values: Sequence[float],
        quantum_kernel: TrainableFidelityQuantumKernel,
        data: np.ndarray,
        labels: np.ndarray,
    ):
        # Bind the user parameter values
        quantum_kernel.assign_training_parameters(parameter_values)
        kernel_matrix = quantum_kernel.evaluate(data)
        return labels.T @ kernel_matrix @ labels

    # Wrap our evaluate method so to produce a callable
    # which maps user_param_values to loss scores.
    def get_variational_callable(
        self,
        quantum_kernel: TrainableFidelityQuantumKernel,
        data: np.ndarray,
        labels: np.ndarray,
    ):
        return partial(
            self.evaluate, quantum_kernel=quantum_kernel, data=data, labels=labels
        )


kernel_loss = CustomKernelLoss().get_variational_callable(qkernel, X_train, y_train)
print(kernel_loss([0.75]))

NUM_GRID_POINTS = 100
loss_values = [kernel_loss([val]) for val in np.linspace(0, 2 * np.pi, NUM_GRID_POINTS)]

plt.rcParams["font.size"] = 15
plt.figure(figsize=(8, 4))
plt.plot(np.linspace(0, 2 * np.pi, NUM_GRID_POINTS), loss_values)
plt.xlabel("θ")
plt.ylabel("Kernel Loss")
plt.show()

def kernel_loss_full(training_param_values, kernel, data, labels):
    kernel.assign_training_parameters(training_param_values)
    kernel_matrix = kernel.evaluate(data)
    return labels.T @ kernel_matrix @ labels


kernel_loss = partial(kernel_loss_full, kernel=qkernel, data=X_train, labels=y_train)
print(kernel_loss([0.75]))

#import qiskit.tools.jupyter
#
#%qiskit_version_table
#%qiskit_copyright