import matplotlib.pyplot as plt
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import execute, Aer
from qiskit.tools.visualization import circuit_drawer
import numpy as np
import time

qr = QuantumRegister(1)
cr = ClassicalRegister(1)

circuit = QuantumCircuit(qr, cr)

circuit.h(qr[0])
circuit.measure(qr[0],cr[0])

backend = Aer.get_backend('qasm_simulator')

exponent = np.arange(2,20,1) # 2^25 iterations is as much as I can calulate in a reasonable amound of time.
expectation = []
for s in [2**n for n in exponent]:
    job = execute(circuit, backend, shots=s)
    result=job.result().get_counts(circuit)
    expectation.append((result['0']-result['1'])/(result['0']+result['1']))

fig = plt.figure(figsize=(10,5))
ax = fig.add_subplot()
ax.scatter(exponent, expectation)
ax.plot(exponent, exponent*0, c=(0.7,.7,0.7))
ax.set_ylabel('<Z> expectation value')
ax.set_xlabel('log2(number of itterations)')

import matplotlib.pyplot as plt
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import execute, Aer
from qiskit.visualization import circuit_drawer
import numpy as np
import pandas as pd
import time

def measure_z(circuit, c_bit, shots=1024):
    """
    Takes a circuit and a classical bit, and returns the expectation value of the Z operator.
    """
    circuit.measure([0],c_bit)
    job = execute(circuit, backend, shots=shots)
    result = job.result().get_counts(circuit) 
    
    keys = list(result.keys())
    if set(keys) == set(['0','1']):
        return (result['0']-result['1'])/(result['0']+result['1'])   
    elif keys == ['0']:
        return 1.0
    elif keys == ['1']:
        return -1.0

def measure_x(circuit, c_bit, shots=1024):
    """
    Takes a circuit and a classical bit, and returns the expectation value of the X operator.
    """
    circuit.h([0])
    
    circuit.measure([0],c_bit)
    job = execute(circuit, backend, shots=shots)
    result = job.result().get_counts(circuit)
    
    keys = list(result.keys())
    if set(keys) == set(['0','1']):
        return (result['0']-result['1'])/(result['0']+result['1'])   
    elif keys == ['0']:
        return 1.0
    elif keys == ['1']:
        return -1.0


def measure_y(circuit, c_bit, shots=1024):
    """
    Takes a circuit and a classical bit, and returns the expectation value of the Y operator.
    """
    circuit.s([0])
    circuit.s([0])
    circuit.s([0]) # using S^-1 = S^3
    circuit.h([0])
    
    circuit.measure([0],c_bit)
    job = execute(circuit, backend, shots=shots)
    result = job.result().get_counts(circuit) 
    
    keys = list(result.keys()) 
    if set(keys) == set(['0','1']):
        return (result['0']-result['1'])/(result['0']+result['1'])   
    elif keys == ['0']:
        return 1.0
    elif keys == ['1']:
        return -1.0

def measure_method_2(theta, circuit, cr, shots=1024):
    """
    Takes a parameter for the matrix M, and an input state circuit with an attached a classical register, and returns <M> 
    via the standard VQE method.
    """
    xcircuit = circuit.copy()
    x = measure_x(xcircuit, cr[0], shots=shots)
    ycircuit = circuit.copy()
    circuit_drawer(ycircuit)
    y = measure_y(ycircuit, cr[0], shots=shots)
    M = np.cos(theta)*x + np.sin(theta)*y
    return M

def measure_method_1(theta, circuit, cr, shots=1024):
    """
    Takes a parameter for the matrix M, and an input state circuit with an attached a classical register, and returns <M> 
    via the unitary basis rotation method.
    """
    zcircuit = circuit.copy()
    zcircuit.p(-theta, [0])
    zcircuit.h([0])
    M = measure_z(zcircuit, cr[0], shots=shots)
    return M


qr = QuantumRegister(1)
cr = ClassicalRegister(1)
backend = Aer.get_backend('qasm_simulator')

shots = 2**0 # Number of shots to obtain expectation value
N = 2**9 # The number of iterations used to see statistics of method

# Method 1

zcircuit = QuantumCircuit(qr, cr)
M1_1 = []
start=time.time()
for i in range(N):
    M1_1.append(measure_method_1(np.pi/4, zcircuit, cr, shots=shots))
t1_1 = time.time()-start

print(f'Method 1: \nTime elapsed for {N} iterations of {shots} shots: ', t1_1)
print('Mean:', np.mean(M1_1))
print('Std. dev.:', np.std(M1_1))

# Method 2

circuit = QuantumCircuit(qr, cr)
M2_1 = []
start=time.time()
for i in range(N):
    M2_1.append(measure_method_2(np.pi/4, circuit, cr, shots=shots))
t2_1 = time.time()-start

print(f'Method 2: \nTime elapsed for {N} iterations of {shots} shots: ', t2_1)
print('Mean:', np.mean(M2_1))
print('Std. dev.:', np.std(M2_1))

fig1, ax1 = plt.subplots(figsize=(8,5))
bins=np.histogram(np.hstack((M1_1,M2_1)), bins=40)[1]
ax1.hist(M1_1, bins, color='b', alpha=0.5, label=f'method 1 ({round(t1/N*1000, 2)} ms/iter)')
ax1.hist(M2_1, bins, color='r', alpha=0.5, label=f'method 2 ({round(t2/N*1000, 2)} ms/iter)')
ax1.axvline(0.0, linestyle='-', c='g', label='True value')
ax1.set_xlabel(r'Estimate of $\langle M\rangle$')
ax1.set_ylabel('Frequency')
ax1.set_title(f'Results of {N} simulations of {shots} shots for '+r'$\theta=\pi/4, \phi=\pi/8$')
ax1.legend()

def init_circuit(phi, qr, cr):
    circuit = QuantumCircuit(qr, cr)
    circuit.h([0])
    circuit.p(-2*phi, [0])
    circuit.h([0]) # Ignoring global phase
    return circuit
    

# Method 1

zcircuit = init_circuit(np.pi/4, qr, cr)
M1_2 = []
start=time.time()
for i in range(N):
    M1_2.append(measure_method_1(np.pi/8, zcircuit, cr, shots=shots))
t1_2 = time.time()-start


print(f'Method 1: \nTime elapsed for {N} iterations of {shots} shots: ', t1_2)
print('Mean:', np.mean(M1_2))
print('Std. dev.:', np.std(M1_2))

# Method 2

circuit = init_circuit(np.pi/4, qr, cr)
M2_2 = []
start=time.time()
for i in range(N):
    M2_2.append(measure_method_2(np.pi/8, circuit, cr, shots=shots))
t2_2 = time.time()-start

print(f'Method 2: \nTime elapsed for {N} iterations of {shots} shots: ', t2_2)
print('Mean:', np.mean(M2_2))
print('Std. dev.:', np.std(M2_2))

fig1, ax1 = plt.subplots(figsize=(8,5))
bins=np.histogram(np.hstack((M1_2,M2_2)), bins=40)[1]
ax1.hist(M1_2, bins, color='b', alpha=0.5, label=f'method 1 ({round(t1/N*1000, 2)} ms/iter)')
ax1.hist(M2_2, bins, color='r', alpha=0.5, label=f'method 2 ({round(t2/N*1000, 2)} ms/iter)')
ax1.axvline(0.3827, linestyle='-', c='g', label='True value')
ax1.set_xlabel(r'Estimate of $\langle M\rangle$')
ax1.set_ylabel('Frequency')
ax1.set_title(f'Results of {N} simulations of {shots} shots for '+r'$\theta=\pi/4, \phi=\pi/8$')
ax1.legend()

shots = np.arange(1000, 100000, 1000)
s_exps  = np.linspace(7, 17, 200)
#shots = [(2**exp)//1 for exp in s_exps]

# Method 1
zcircuit = init_circuit(np.pi/4, qr, cr)
M1 = []
t1 = []
for s in shots:
    start=time.time()
    M1.append(measure_method_1(np.pi/8, zcircuit, cr, shots=s))
    end=time.time()
    t1.append(end-start)
    
# Method 2
circuit = init_circuit(np.pi/4, qr, cr)
M2 = []
t2 = []
for s in shots:
    start=time.time()
    M2.append(measure_method_2(np.pi/8, circuit, cr, shots=s))
    end=time.time()
    t2.append(end-start)
    


fig, axs = plt.subplots(1,2,figsize=(16,5))
axs[0].plot(shots,t1, label='method 1')
axs[0].plot(shots,t2, label='method 2')
axs[0].set_xlabel('Number of shots')
axs[0].set_ylabel('Time taken')
axs[0].legend()

axs[1].scatter(shots,M1, label='method 1')
axs[1].scatter(shots,M2, label='method 2')
axs[1].axhline(np.sin(np.pi/8), label='expected value', c=(.7,.7,.7))
#axs[1].set_xlabel(r'$log_2(shots)$')
axs[1].set_ylabel(r'$\langle M\rangle')
axs[1].legend()

tps_1 = np.mean(np.array(t1)/np.array(shots))
tps_2 = np.mean(np.array(t2)/np.array(shots))
print('Ratio of time taken (method 2:method 1) = ', round(tps_2/tps_1, 3))

#s_exps  = np.linspace(7, 14, 10)
#shots = [(2**exp)//1 for exp in s_exps]
shots_std = np.linspace(500,2000,15)
N=50

# Method 1
zcircuit = init_circuit(np.pi/4, qr, cr)
std1 = []
for s in shots_std:
    m = []
    for i in range(N):
        m.append(measure_method_1(np.pi/8, zcircuit, cr, shots=s))
    std1.append(np.std(m))
    
# Method 2
circuit = init_circuit(np.pi/4, qr, cr)
std2=[]
for s in shots_std:
    m=[]
    for i in range(N):
        m.append(measure_method_2(np.pi/8, circuit, cr, shots=s))
    std2.append(np.std(m))

        

fig, axs = plt.subplots(1,2,figsize=(16,5))
axs[0].plot(shots,t1, label='method 1', c='r')
axs[0].plot(shots,t2, label='method 2', c='b')
axs[0].set_xlabel('Number of shots')
axs[0].set_ylabel('Time taken')
axs[0].legend()

rt_s = np.sqrt(1/np.array(shots_std))
axs[1].scatter(rt_s, std1, label='method 1', c='r')
axs[1].plot(rt_s, rt_s*np.polyfit(rt_s, std1,1)[0]+ np.polyfit(rt_s, std1,1)[1],
            c='r', alpha=.5, label='m='+str(round(np.polyfit(rt_s, std1,1)[0],3)))

axs[1].scatter(rt_s, std2, label='method 2', c='b')
axs[1].plot(rt_s, rt_s*np.polyfit(rt_s, std2,1)[0]+ np.polyfit(rt_s, std2,1)[1],
            c='b', alpha=.5, label='m='+str(round(np.polyfit(rt_s, std2,1)[0],3)))
axs[1].set_xlabel(r'shots$^{-1/2}$')
axs[1].set_ylabel(r'$\Delta\langle M\rangle$')
axs[1].legend()

