from qiskit import * 
import matplotlib.pyplot as plt 
from qiskit.extensions import UnitaryGate
from qiskit.circuit import add_control 
from qiskit.tools.visualization import plot_bloch_multivector,plot_histogram
import numpy as np 
import sys  
sys.path.append("..")

from Modules.vanilla_qpe import QPE 

U= UnitaryGate(data = np.array([[1,0],
              [0,np.exp(2*np.pi*1j*(1/5))]]))


qpe_circ1 = QPE(precision=4,unitary=U).get_QPE(show=True,save = False) 

q = QuantumCircuit(5,4)
q.x(4)
q.barrier()

q = q.compose(qpe_circ1,qubits = [0,1,2,3,4])
q.draw('mpl')

q.measure([0,1,2,3],[0,1,2,3])
q.draw('mpl')

count = execute(q,backend=Aer.get_backend('qasm_simulator')).result().get_counts()
fig = plot_histogram(count,title = "Phase estimates for $\\theta = 1/5$")
fig.savefig("Phase 0.2 estimate.JPG",dpi = 200)
fig 

q = QuantumCircuit(3,name = 'Unitary')
q.cp(2*np.pi*(1/5),1,2)
q.draw('mpl')

qpe_circ2 = QPE(precision=4,unitary=q).get_QPE(show=True,save=True)

q = QuantumCircuit(7,4)
q.x([5,6])
q.barrier()
q.append(qpe_circ2, qargs = [0,1,2,3,4,5,6])
q.measure([0,1,2,3],[0,1,2,3])
q.draw('mpl')

count = execute(q,backend=Aer.get_backend('qasm_simulator')).result().get_counts()
plot_histogram(count)

q = QuantumCircuit(7,5)
q.x(4)
q.compose(qpe_circ1,qubits = [0,1,2,3,4], inplace = True)
# q = q.compose(qpe_circ,qubits = [0,1,2,3,4])
q.draw('mpl')

q.measure([0,1,2,3],[0,1,2,3])
q.draw('mpl')

count = execute(q,backend=Aer.get_backend('qasm_simulator')).result().get_counts()
plot_histogram(count)