from typing import List, Callable, Union

from qiskit import QuantumCircuit
from qiskit.circuit import ParameterVector

# To visualize circuit creation process
from qiskit.visualization import circuit_drawer

# For a dataset with 12 features; and 2 features per qubit
FEATURE_DIMENSION = 12
NUM_QUBITS = int(FEATURE_DIMENSION / 2)

# Qiskit feature maps should generally be QuantumCircuits or extensions of QuantumCircuit
feature_map = QuantumCircuit(NUM_QUBITS)
user_params = ParameterVector("θ", NUM_QUBITS)

# Create circuit layer with trainable parameters
for i in range(NUM_QUBITS):
    feature_map.ry(user_params[i], feature_map.qubits[i])

print(circuit_drawer(feature_map))

# Linear entanglement
entanglement = [[i, i + 1] for i in range(NUM_QUBITS - 1)]

for source, target in entanglement:
    feature_map.cz(feature_map.qubits[source], feature_map.qubits[target])

feature_map.barrier()

print(circuit_drawer(feature_map))

input_params = ParameterVector("x", FEATURE_DIMENSION)
for i in range(NUM_QUBITS):
    feature_map.rz(input_params[2 * i + 1], feature_map.qubits[i])
    feature_map.rx(input_params[2 * i], feature_map.qubits[i])

print(circuit_drawer(feature_map))

class ExampleFeatureMap(QuantumCircuit):
    """The Example Feature Map circuit"""

    def __init__(
        self,
        feature_dimension: int,
        entanglement: Union[str, List[List[int]], Callable[[int], List[int]]] = None,
        name: str = "ExampleFeatureMap",
    ) -> None:
        """Create a new Example Feature Map circuit.
        Args:
            feature_dimension: The number of features
            entanglement: Entanglement scheme to be used in second layer
            name: Name of QuantumCircuit object

        Raises:
            ValueError: ExampleFeatureMap requires an even number of input features
        """
        if (feature_dimension % 2) != 0:
            raise ValueError(
                """
            Example feature map requires an even number of input features.
                """
            )
        self.feature_dimension = feature_dimension
        self.entanglement = entanglement
        self.training_parameters = None

        # Call the QuantumCircuit initialization
        num_qubits = feature_dimension / 2
        super().__init__(
            num_qubits,
            name=name,
        )

        # Build the feature map circuit
        self._generate_feature_map()

    def _generate_feature_map(self):
        # If no entanglement scheme specified, use linear entanglement
        if self.entanglement is None:
            self.entanglement = [[i, i + 1] for i in range(self.num_qubits - 1)]

        # Vector of data parameters
        input_params = ParameterVector("x", self.feature_dimension)

        training_params = ParameterVector("θ", self.num_qubits)
        # Create an initial rotation layer of trainable parameters
        for i in range(self.num_qubits):
            self.ry(training_params[i], self.qubits[i])

        self.training_parameters = training_params

        # Create the entanglement layer
        for source, target in self.entanglement:
            self.cz(self.qubits[source], self.qubits[target])

        self.barrier()

        # Create a circuit representation of the data group
        for i in range(self.num_qubits):
            self.rz(input_params[2 * i + 1], self.qubits[i])
            self.rx(input_params[2 * i], self.qubits[i])

feature_map = ExampleFeatureMap(feature_dimension=10)
circuit_drawer(feature_map)

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