import numpy as np
import math
from qiskit import QuantumCircuit
from qiskit.circuit.library import GroverOperator, MCMT, ZGate
from qiskit.visualization import plot_distribution
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
backend = GenericBackendV2(num_qubits=5)

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(i,zero_inds)
        qc.ry(j,zero_inds)
        qc.rz(k,zero_inds)
            
        qc.compose(MCMT(ZGate(), num_qubits - 1, 1), inplace=True)
        qc.x(zero_inds)
    return qc

def printToConsole():
    # Format counts data
    formatted_output = ""
    #for state in ['000', '001', '010', '011', '100', '101', '110', '111']:
    for state in ['0000', '0001', '0010', '0011', '0100', '0101', '0110', '0111', 
                '1000', '1001', '1010', '1011', '1100', '1101', '1110', '1111']:
        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)
        file.write(f"{x} {y} {z}\n")
    if(printDebug):
        import matplotlib.pyplot as plt
            
        # Parse counts data
        counts_data = {state: counts.get(state, 0) for state in ['0000', '0001', '0010', '0011', '0100', '0101', '0110', '0111', 
                              '1000', '1001', '1010', '1011', '1100', '1101', '1110', '1111']}
        #for state in ['000', '001', '010', '011', '100', '101', '110', '111']}
            
        # Plot counts
        plt.bar(counts_data.keys(), counts_data.values())
        plt.xlabel('States')
        plt.ylabel('Counts')
        plt.title('Quantum Circuit Measurement Results')
        plt.show()
        print(counts_data)

marked_states = ["0010", "1000"]
printDebug = False
randomXYZ = False

#RX, RY, RZ trasformators causing error in the grover oracle by rotating state by v/100
fromRot = -100
toRot = 101
XYZ = 2 #1->RX, 2->RY, 3->RZ

num_simulations = 10 #averaging this many simulations, (good for graphing)

for i in range(fromRot, toRot,1):
    # Initialize empty dictionary for averages
    averages = {}
    x=0
    y=0
    z=0
    # Loop for multiple simulations
    for _ in range(num_simulations):
        zaj=i
        if(randomXYZ):
            k = random.randrange(-100,100)#-1 - +1 fok kozott random zaj 

        if XYZ == 1:
            x=zaj/100
        elif XYZ == 2:
            y=zaj/100
        elif XYZ == 3:
            z=zaj/100
            
        oracle = grover_oracle(marked_states, x, y, z)
            
        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)
        qc.h(range(grover_op.num_qubits))
        qc.compose(grover_op.power(optimal_num_iterations), inplace=True)
        qc.measure_all()
        backend = GenericBackendV2(num_qubits=4)
        transpiled_circuit = transpile(qc, backend)
        job = backend.run(transpiled_circuit)
        counts = job.result().get_counts()
                    
        # Update averages for each state
        #for state in ['000', '001', '010', '011', '100', '101', '110', '111']:
        for state in ['0000', '0001', '0010', '0011', '0100', '0101', '0110', '0111', 
                        '1000', '1001', '1010', '1011', '1100', '1101', '1110', '1111']:
            averages[state] = averages.get(state, 0) + counts.get(state, 0)
                
    # Calculate averages
    for state in averages:
        averages[state] /= num_simulations
                
    # Write averages to a text file
    with open('averages_data.txt', 'a') as file:
        for state in averages:
            file.write(f"{int(averages[state])} ")
        file.write(f"{x} {y} {z}")
        file.write("\n")

if(printDebug):
    printToConsole()
        
    
print("All done")