from qiskit import QuantumCircuit, execute, transpile, Aer 
from qiskit.extensions import UnitaryGate,Initialize
from qiskit.quantum_info import Statevector 
from qiskit.tools.visualization import plot_bloch_vector
from qiskit.tools.visualization import plot_histogram,plot_bloch_multivector  
import numpy as np 
import os  
from time import sleep
import sys
sys.path.append("..")
from scipy.stats import unitary_group 
import matplotlib.pyplot as plt 
%matplotlib inline 

os.makedirs(name = "Plots/1-qubit/",exist_ok=True)

from Modules.vanilla_qpe import QPE 
from Modules.iterative_qpe import IQPE
from Modules.normal_SPEA import SPEA 
from Modules.kitaev_qpe import KQPE

# x gate 
u1 = QuantumCircuit(1)
u1.x(0)
u1.draw('mpl')

# x gate 
u2 = QuantumCircuit(1)
u2.p(2*np.pi*(1/6),0)
u2.draw('mpl')

backend = Aer.get_backend('qasm_simulator')

os.makedirs(name = "Plots/1-qubit/x_gate/",exist_ok=True)
x_path = "Plots/1-qubit/x_gate/"

estimates = []
for precision in range(1,9): 
    qpe = QPE(precision=precision,unitary = u1)
    qpe_circ = qpe.get_QPE(show = False)
    qc = QuantumCircuit(precision+1,precision)
    qc.x(precision)
    qc.h(precision)
    qc.append(qpe_circ, qargs = list(range(precision+1)))
    qc.measure(list(range(precision)),list(range(precision)))
#     display(qc.draw('mpl')) 
    counts = execute(qc,backend=backend,shots = 2**10).result().get_counts()
    m = -1 
    for i,j in counts.items():
        if j > m:
            m = j 
            phase = i 
    factor = 0.5
    ans = 0.0 
    for i in range(precision):
        ans+= int(phase[i])*factor 
        factor/=2 
    estimates.append(ans)

plt.title("Phase by basic QPE for the unitary X",fontsize = 16)
plt.plot([i for i in range(1,9)],estimates,color = 'g',linewidth = 2,marker = 'o',label = 'Estimates')
plt.xlabel('Precision')
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(x_path+"basic_qpe.JPG",dpi = 200)

estimates = []
for precision in range(1,9): 
    qpe = IQPE(precision=precision,unitary = u1,unknown=True)
    qc = QuantumCircuit(precision+1,precision)
    qc.x(1)
    qc.h(1)    
    phase = qpe.get_circuit_phase(qc, qubits = [1], ancilla = 0,clbits = list(range(precision)),show = False)
#     display(qc.draw('mpl')) 
    estimates.append(phase[1])

plt.title("Phase by iterative QPE for the unitary X",fontsize = 16)
plt.plot([i for i in range(1,9)],estimates,color = 'g',linewidth = 2,marker = 'o',label = 'Estimates')
plt.xlabel('Precision')
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(x_path+"iqpe.JPG",dpi = 200)

estimates = []
for resolution in range(10,40,5): 
    qpe = SPEA(resolution = resolution,unitary = u1, error =3)
    result = qpe.get_eigen_pair(backend = backend)
#     display(qc.draw('mpl')) 
    print("Result : ",result)
    estimates.append(result['theta'])

plt.title("Phase by SQPE for the unitary X",fontsize = 16)
plt.plot([i for i in range(10,40,5)],estimates,color = 'g',linewidth = 2,marker = 'o',label = 'Estimates')
plt.plot([5,50],[0.5,0.5],color = 'black')
plt.plot([5,50],[0,0],color = 'black')
plt.plot([5,50],[1,1],color = 'black',label = 'Correct')
plt.xlabel('Precision')
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(x_path+"stat_qpe.JPG",dpi = 200)

estimates = []
for precision in range(5,16): 
    kqpe = KQPE(precision=precision,unitary = u1)
    kqpe_circ = kqpe.get_circuit()
    qc = QuantumCircuit(2+1,2)
    qc.x(2)
    qc.h(2)
    qc.append(kqpe_circ,qargs = [0,1,2])
    phase = kqpe.get_phase(qc,backend = backend, ancilla = [0,1],clbits = [0,1],show = False)
    print("Precision :",precision)
    print("Phase:",phase)
    estimates.append(phase[0])

plt.title("Phase by Kitaev's Algorithm for the X gate",fontsize = 16)
plt.plot([i for i in range(5,16)],estimates,color = 'magenta',alpha = 0.6,linewidth = 2,marker = 's',label = 'Estimates')
plt.plot([5,17],[0.5,0.5],color = 'black')
plt.plot([5,17],[-0.5,-0.5],color = 'black')
plt.xlabel('Precision')     
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(x_path+"kitaev_qpe.JPG",dpi = 200)

os.makedirs(name = "Plots/1-qubit/phase_gate/",exist_ok=True)
phase_path = "Plots/1-qubit/phase_gate/"

estimates = []
for precision in range(1,9): 
    qpe = QPE(precision=precision,unitary = u2)
    qpe_circ = qpe.get_QPE(show = False)
    qc = QuantumCircuit(precision+1,precision)
    qc.x(precision)
    qc.append(qpe_circ, qargs = list(range(precision+1)))
    qc.measure(list(range(precision)),list(range(precision)))
#     display(qc.draw('mpl')) 
    counts = execute(qc,backend=backend,shots = 2**10).result().get_counts()
    m = -1 
    for i,j in counts.items():
        if j > m:
            m = j 
            phase = i 
    factor = 0.5
    ans = 0.0 
    for i in range(precision):
        ans+= int(phase[i])*factor 
        factor/=2 
    estimates.append(ans)

plt.title("Phase by basic QPE for the unitary with phase(1/6)",fontsize = 16)
plt.plot([i for i in range(1,9)],estimates,color = 'red',alpha = 0.8, linewidth = 2,marker = 'o',label = 'Estimates')
plt.plot([0,10],[0.1667,0.1667], color = 'black')
plt.xlabel('Precision')
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(phase_path+"basic_qpe.JPG",dpi = 200)

estimates = []
for precision in range(1,9): 
    qpe = IQPE(precision=precision,unitary = u2,unknown=True)
    qc = QuantumCircuit(1+1,precision)
    qc.x(1)   
    phase = qpe.get_circuit_phase(qc, qubits = [1], ancilla = 0,clbits = list(range(precision)),show = False)
#     display(qc.draw('mpl')) 
    estimates.append(phase[1])

plt.title("Phase by iterative QPE for the unitary with phase(1/6)",fontsize = 16)
plt.plot([i for i in range(1,9)],estimates,color = 'orange',alpha = 0.7,linewidth = 2,marker = 'o',label = 'Estimates')
plt.plot([0,10],[0.1667,0.1667], color = 'black')
plt.xlabel('Precision')
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(phase_path+"iqpe.JPG",dpi = 200)

estimates = []
for resolution in range(10,40,5): 
    qpe = SPEA(resolution = resolution,unitary = u2, error =4)
    result = qpe.get_eigen_pair(backend = backend)
    print("Result : ",result)
    estimates.append(result['theta'])

plt.title("Phase by SQPE for the unitary with phase(1/6)",fontsize = 16)
plt.plot([i for i in range(10,40,5)],estimates,color = 'g',alpha = 0.7,linewidth = 2,marker = 'o',label = 'Estimates')
plt.plot([5,50],[0.1667,0.1667],color = 'black')
plt.plot([5,50],[0,0],color = 'black')
plt.plot([5,50],[1,1],color = 'black',label = 'Correct')
plt.xlabel('Precision')     
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(phase_path+"stat_qpe.JPG",dpi = 200)

estimates = []
for precision in range(5,16): 
    kqpe = KQPE(precision=precision,unitary = u2)
    kqpe_circ = kqpe.get_circuit()
    qc = QuantumCircuit(2+1,2)
    qc.x(2)   
    qc.append(kqpe_circ,qargs = [0,1,2])
    phase = kqpe.get_phase(qc,backend = backend, ancilla = [0,1],clbits = [0,1],show = False)
#     display(qc.draw('mpl')) 
    print("Precision :",precision)
    print("Phase:",phase)
    estimates.append(phase[0])

plt.title("Phase by Kitaev's Algorithm for the unitary with phase(1/6)",fontsize = 16)
plt.plot([i for i in range(5,16)],estimates,color = 'magenta',alpha = 0.6,linewidth = 2,marker = 's',label = 'Estimates')
plt.plot([5,17],[0.1667,0.1667],color = 'black')
plt.xlabel('Precision')     
plt.ylabel('Estimates')
plt.grid()
plt.legend()
plt.savefig(phase_path+"kitaev_qpe.JPG",dpi = 200)

