from qiskit import QuantumCircuit, Aer, transpile, assemble
from qiskit.aqua.components.optimizers import COBYLA
from qiskit.circuit import Parameter
import numpy as np

# Define the desired probability distribution
target_probs = [0.1666666666] * 6

# Define the quantum circuit
circ = QuantumCircuit(6, 6)
circ.h(range(6))

# Add the parameterized Ansatz using Parameter
params = [Parameter(f'var{i}') for i in range(6)]
circ.rx(params[0], 0)
circ.rx(params[1], 1)
circ.rx(params[2], 2)
circ.rx(params[3], 3)
circ.rx(params[4], 4)
circ.rx(params[5], 5)


# Define the optimization objective function
def objective_function(params):
    # Update the circuit with new parameters
    updated_circ = circ.bind_parameters({params[i]: optimal_params[i] for i in range(6)})

    # Compile and run the circuit
    t_circ = transpile(updated_circ, Aer.get_backend('qasm_simulator'))
    qobj = assemble(t_circ)
    result = Aer.get_backend('qasm_simulator').run(qobj).result()

    # Calculate the probability distribution
    counts = result.get_counts()
    probs = [counts.get(format(i, '06b'), 0) / result.shots for i in range(2 ** 6)]

    # Calculate the objective function (sum of squared differences)
    return sum((probs[i] - target_probs[i]) ** 2 for i in range(2 ** 6))


# Optimize the circuit parameters
optimizer = COBYLA(maxiter=1000)
optimal_params = optimizer.optimize(num_vars=6, objective_function=objective_function)

# Print the optimal parameters
print("Optimal Parameters:", optimal_params)