#!/usr/bin/env python
# coding: utf-8

# # I-Python/Jupyter Notebook: Demo for PHY682 QIS
# ### Created by Tzu-Chieh Wei <br> Materials collected and revised from various places: QISKIT Summer School 2020, QISKIT Book, etc.  [Updated Sep 2021]
# ### Simple VQE illustration

# In[1]:


# import packages
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister


# In[2]:


# first variational form and Z measurement
def get_var_formZ(params):
    qr = QuantumRegister(1, name="q")
    cr = ClassicalRegister(1, name='c')
    qc = QuantumCircuit(qr, cr)
    qc.u(params[0], params[1], params[2], qr[0])
    qc.measure(qr,cr)
    return qc


# In[3]:


# variational form and X measurement
def get_var_formX(params):
    qr = QuantumRegister(1, name="q")
    cr = ClassicalRegister(1, name='c')
    qc = QuantumCircuit(qr, cr)
    qc.u(params[0], params[1], params[2], qr[0])
    qc.h(qr[0])
    qc.measure(qr,cr)
    return qc


# In[4]:


from qiskit.circuit import Parameter
theta1 = Parameter('θ1')
theta2 = Parameter('θ2')
theta3 = Parameter('θ3')
ex_qc=get_var_formX([theta1,theta2,theta3])
ex_qc.draw('mpl')


# In[5]:


ex_qc=get_var_formZ([theta1,theta2,theta3])
ex_qc.draw('mpl')


# In[6]:


# proability distribution P(0) and P(1)
def get_probability_distribution(counts):
    mysum=sum([v for  v in counts.values()])
    keys=['0','1']
    output_distr =[]
    for key in keys:
        if key not in counts.keys():
            output_distr.append(0)
        else:
            output_distr.append(counts[key]/mysum)
    return output_distr


# In[7]:


# choose backend to be the simulator
from qiskit import Aer, execute
backend = Aer.get_backend("qasm_simulator")
NUM_SHOTS = 10000


# In[8]:


# define energy for Hamiltonian - B[0] sigma_z - B[1] sigma_x
# B is a list of two objects B=[Bz, Bx]
def energy(output_distrZ,output_distrX,B):
    cost = -B[0]*(output_distrZ[0] -output_distrZ[1])-B[1]* (output_distrX[0] -output_distrX[1])
    return cost


# In[9]:


# the objective function is to evaluate the energy cost by running the quantum circuits
def objective_function(params,B):
    # Obtain quantum circuit instances from the paramters
    qc1 = get_var_formZ(params)
    qc2 = get_var_formX(params)
    # Execute the quantum circuit to obtain the probability distribution associated with the current parameters
    runjob = execute([qc1, qc2],backend,shots=10000)
    results=runjob.result()
    
    # Obtain the counts for each measured state, and convert those counts into a probability vector
    output_distr1 = get_probability_distribution(results.get_counts(qc1))
    output_distr2 = get_probability_distribution(results.get_counts(qc2))
    # Calculate the cost as the distance between the output distribution and the target distribution
    cost = energy(output_distr1,output_distr2,B)
    return cost


# In[10]:


import numpy as np
from qiskit.algorithms.optimizers import COBYLA

nshots=10000
# Initialize the COBYLA optimizer via Qiskit
optimizer = COBYLA(maxiter=500, disp=True,tol=1e-6)

Bfield =np.random.rand(2)
print('(Bz,Bx)=',Bfield)

# Create the initial parameters (noting that our single qubit variational form has 3 parameters)
params = np.random.rand(3)


def myobject(params):
    fvalue=objective_function(params,Bfield)
    return fvalue

counts = []
values = []
myparams = []
def store_intermediate_result(eval_count, parameters, mean):
        counts.append(eval_count)
        values.append(mean)
        myparams.append(parameters)
        
ret = optimizer.optimize(num_vars=3, objective_function=myobject, initial_point=params)

# 
qc1 = get_var_formZ(ret[0])
qc2 = get_var_formX(ret[0])
#from qiskit.visualization import circuit_drawer
#circuit_drawer(qc1)
#circuit_drawer(qc2)  cannot draw using the returned circuit!!
counts1 = execute(qc1, backend, shots=nshots).result().get_counts(qc1)
counts2 = execute(qc2, backend, shots=nshots).result().get_counts(qc2)
dst1=get_probability_distribution(counts1)
dst2=get_probability_distribution(counts2)


myEnergy=energy(dst1,dst2,Bfield)
exactEnergy=-np.sqrt(Bfield[0]**2+Bfield[1]**2)
print('Num of shots for each circuit=',nshots)
print('Estimated minimum energy=',myEnergy,' vs. exact=',exactEnergy,' error=',myEnergy-exactEnergy)
print("Sigma_z:", dst1)
print("Sigma_x:", dst2)
print("Parameters Found:", ret[0])


# In[11]:


print(ret[0],ret[1],ret[2])


# In[12]:


from scipy.optimize import minimize


# In[13]:


Bfield =np.random.rand(2)
print('(Bz,Bx)=',Bfield)
counts = []
values = []
myparams = []
def store_intermediate_result(parameters):
        print(parameters)
        myparams.append(parameters)
        
params = np.random.rand(3)     

#Using directly SciPy optimizer
#ret2=minimize(myobject,params,method='COBYLA',options={'maxiter':500, 'disp':True,'tol':1e-6},callback=store_intermediate_result)
ret2=minimize(myobject,params,method='L-BFGS-B',options={'maxiter':500, 'disp':True,'maxcor': 10, 'ftol': 2.220446049250313e-09, 'gtol': 1e-05, 'eps': 1e-08, 'maxfun': 15000, 'maxiter': 15000, 'iprint': - 1, 'maxls': 20, 'finite_diff_rel_step': None},callback=store_intermediate_result)


# In[14]:


print(ret2)
exactEnergy=-np.sqrt(Bfield[0]**2+Bfield[1]**2)
print('Exact energy= ',exactEnergy)


# In[16]:


params = np.random.rand(3)        
ret2=minimize(myobject,params,method='COBYLA',options={'maxiter':500, 'disp':True,'tol':1e-6},callback=store_intermediate_result)
#ret2=minimize(myobject,params,method='L-BFGS-B',options={'maxiter':500, 'disp':True,'maxcor': 10, 'ftol': 2.220446049250313e-09, 'gtol': 1e-05, 'eps': 1e-08, 'maxfun': 15000, 'maxiter': 15000, 'iprint': - 1, 'maxls': 20, 'finite_diff_rel_step': None},callback=store_intermediate_result)


# In[17]:


print(ret2)
exactEnergy=-np.sqrt(Bfield[0]**2+Bfield[1]**2)
print(exactEnergy)


# In[18]:


ret2=minimize(myobject,ret2.x,method='L-BFGS-B',options={'maxiter':500, 'disp':True,'maxcor': 10, 'ftol': 2.220446049250313e-09, 'gtol': 1e-05, 'eps': 1e-08, 'maxfun': 15000, 'maxiter': 15000, 'iprint': - 1, 'maxls': 20},callback=store_intermediate_result)


# In[19]:


print(ret2)
exactEnergy=-np.sqrt(Bfield[0]**2+Bfield[1]**2)
print(exactEnergy)


# ## Qiskit has built-in packages

# In[15]:


from qiskit.aqua.operators import WeightedPauliOperator # deprecated
from qiskit.opflow import X, Z, I
from qiskit.algorithms import NumPyEigensolver
#from qiskit.aqua.components.variational_forms import RY, RYRZ, SwapRZ   #deprecated
from qiskit.algorithms import VQE
from qiskit.circuit.library import TwoLocal
from qiskit.algorithms.optimizers import COBYLA, L_BFGS_B, SLSQP, SPSA


# In[16]:


def hamiltonian(a, b, c, d):
    """
    Creates a*I + b*Z + c*X + d*Y pauli sum 
    that will be our Hamiltonian operator.
    
    """
    pauli_dict = {
        'paulis': [{"coeff": {"imag": 0.0, "real": a}, "label": "I"},
                   {"coeff": {"imag": 0.0, "real": b}, "label": "Z"},
                   {"coeff": {"imag": 0.0, "real": c}, "label": "X"},
                   {"coeff": {"imag": 0.0, "real": d}, "label": "Y"}
                   ]
    }
    return WeightedPauliOperator.from_dict(pauli_dict)


def BzBxH(B):
    return B[0]*Z + B[1]*X 


# In[17]:


H=BzBxH(Bfield)


# In[18]:


def store_intermediate_result(eval_count, parameters, mean, std):
            counts.append(eval_count)
            values.append(mean)
            params.append(parameters)
            deviation.append(std)


# In[19]:


depth=1
#var_name='RYRZ'
#var_form = RYRZ(H.num_qubits, depth=depth) # deprecated
var_form=TwoLocal(H.num_qubits,rotation_blocks=['ry','rz'],entanglement=None,reps=0)
print(var_form)
optimizer = COBYLA(maxiter=500, disp=True,tol=1e-6)
counts=[]
values=[]
params=[]
deviation=[]

vqe = VQE(var_form, optimizer, quantum_instance=backend, callback=store_intermediate_result)


# In[20]:


#backend = BasicAer.get_backend('statevector_simulator')
vqe_result = vqe.compute_minimum_eigenvalue(H)
print(vqe_result)


# In[21]:


exact_result = NumPyEigensolver().compute_eigenvalues(operator=H)
exact_energy = min(np.real(exact_result.eigenvalues))


# In[23]:


print('VQE result: ',vqe_result.optimal_value,' vs Exact: ',exact_energy)


# In[24]:


# It allows to use callback to keep track of convergence
import matplotlib.pyplot as plt 

plt.plot(counts, abs(exact_energy - values))
plt.xlabel('Eval count')
plt.ylabel('Energy difference from solution reference value')
plt.title('Energy convergence')
plt.yscale('log')
plt.legend(loc='upper right')


# In[25]:


import qiskit
qiskit.__qiskit_version__


# In[ ]: