# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, transpile
from qiskit_ibm_provider import IBMProvider
import qiskit_ibm_provider.jupyter
 
#provider = IBMProvider('ibm-q')
#backend = provider.get_backend('ibmq_vigo')
from qiskit.visualization import *
from ibm_quantum_widgets import *

# qiskit-ibmq-provider has been deprecated.
# Please see the Migration Guides in https://ibm.biz/provider_migration_guide for more detail.
from qiskit_ibm_runtime import QiskitRuntimeService, Sampler, Estimator, Session, Options

# Loading your IBM Quantum account(s)
service = QiskitRuntimeService(channel="ibm_quantum")

# Invoke a primitive. For more details see https://qiskit.org/documentation/partners/qiskit_ibm_runtime/tutorials.html
# result = Sampler("ibmq_qasm_simulator").run(circuits).result()# Built-in modules
import math

# Imports from Qiskit
from qiskit import QuantumCircuit
from qiskit.circuit.library import GroverOperator, MCMT, ZGate
from qiskit.visualization import plot_distribution

# Imports from Qiskit Runtime
from qiskit_ibm_runtime import QiskitRuntimeService, Sampler, Session

from qiskit.providers.basic_provider import BasicProvider
from qiskit.visualization import plot_histogram
from qiskit import transpile
from qiskit.providers.fake_provider import GenericBackendV2

def grover_oracle(marked_states, i, j, k):
    if not isinstance(marked_states, list):
        marked_states = [marked_states]
    num_qubits = len(marked_states[0])
    qc = QuantumCircuit(num_qubits)
    for target in marked_states:
        rev_target = target[::-1]
        zero_inds = [ind for ind in range(num_qubits) if rev_target.startswith("0", ind)]
        qc.x(zero_inds)
        
        """ZAJ 1"""
        qc.rx(0.1*i,zero_inds)
        qc.ry(0.1*j,zero_inds)
        qc.rz(0.1*k,zero_inds)
        
        qc.compose(MCMT(ZGate(), num_qubits - 1, 1), inplace=True)
        qc.x(zero_inds)
    return qc

marked_states = ["011", "100"]
printToConsole = False

for i in range(0,10):
    for j in range(0,10):
        for k in range(0,10):
            oracle = grover_oracle(marked_states, i, j, k)
            grover_op = GroverOperator(oracle)
            optimal_num_iterations = math.floor(
                math.pi / 4 * math.sqrt(2**grover_op.num_qubits / len(marked_states))
            )
            qc = QuantumCircuit(grover_op.num_qubits)
            # Create even superposition of all basis states
            qc.h(range(grover_op.num_qubits))

            # Apply Grover operator the optimal number of times
            qc.compose(grover_op.power(optimal_num_iterations), inplace=True)

            """ZAJ 2"""
            #qc.x(range(grover_op.num_qubits))

            # Measure all qubits
            qc.measure_all()
            #qc.draw('mpl')
            grover_op.decompose().draw("mpl")

            backend = GenericBackendV2(num_qubits=5)

            transpiled_circuit = transpile(qc, backend)

            job = backend.run(transpiled_circuit)
            counts = job.result().get_counts()

            # Format counts data
            formatted_output = ""
            for state in ['000', '001', '010', '011', '100', '101', '110', '111']:
                formatted_output += f"{counts.get(state, 0)} "

            # Write counts data to a text file
            with open('counts_data.txt', 'a') as file:
                file.write(formatted_output + "\n")
            
            
            """# Write counts data to a text file
            with open('counts_data.txt', 'a') as file:
                file.write(f"i={i}j={j}k={k}")
                file.write("State,Count\n")
                for state, count in counts.items():
                    file.write(f"{state},{count}\n")
            if(printToConsole):
                # Print counts data
                print(counts)
                print("State,Count\n")
                for state, count in counts.items():
                    print(f"{state},{count}\n")
                #plot_histogram(counts)"""
print("All done")

