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

# # Variational quantum eigensolving
# 
# In this exercise, we exemplify how the VQE method can be used to find the approximate ground state (GS) of a given Hamiltonian. Here, we take a Heisenberg Hamiltonian and a random Hamiltonian as concrete examples.
# 
# ## Defining the Hamiltonian
# 
# ### The Heisenberg Hamiltonian
# 
# We start by defining the Hamiltonian we want to find an approximate ground state of. Let us start with the Heisenberg Hamiltonian:
# 
# $$H=J\left(\hat{S}_{1}^{x}\hat{S}_{2}^{x}+\hat{S}_{1}^{y}\hat{S}_{2}^{y}+\hat{S}_{1}^{z}\hat{S}_{2}^{z}\right).$$
# 
# In myQLM, Hamiltonians are implemented with ``Observable`` objects (https://myqlm.github.io/qat-core.html#module-qat.core.wrappers.observable).
# 
# 
# 

# In[ ]:


# only if using Google Colab:
get_ipython().system('pip install myqlm')


# In[ ]:


import numpy as np
from qat.core import Observable, Term

nqbits = 2
terms = #...

hamiltonian = Observable(nqbits, pauli_terms=terms)
print("H:", hamiltonian)


# Because $H$ is defined on a small number of qubits, we can afford to diagonalize it exactly and thus compute the exact GS energy:

# In[ ]:


from utils_tuto import make_matrix
H_mat = make_matrix(hamiltonian)
eigvals = np.linalg.eigvalsh(H_mat)
E0 = min(eigvals)


# What is the size of the matrix representation of $H$? For a number of qubits $n$, what is the size in general?
# 
# What does the ground state of $H$ look like? What is special about it?

# ### A slighly more complicated Hamiltonian
# 
# For later purposes, we define a random Hamiltonian with 5 terms acting on 3 qubits. (Do not execute the cell if you want to work with the Heisenberg Hamiltonian first)

# In[ ]:


def make_random_hamiltonian(nqbits, nterms):
    terms = []
    for _ in range(nterms):
        coeff = np.random.random()
        ops = "".join(np.random.choice(["X", "Z"], size=nqbits))
        qbits = np.random.choice(nqbits, size=nqbits, replace=False)
        terms.append(Term(coefficient=coeff, pauli_op=ops, qbits=qbits))
    hamiltonian = Observable(nqbits, pauli_terms=terms)
    return hamiltonian


nqbits = 3
nterms = 5
np.random.seed(1423543) #fixing seed to have reproducible results
hamiltonian_random = make_random_hamiltonian(nqbits, nterms)
print("H:", hamiltonian_random)

H_mat = make_matrix(hamiltonian_random)
eigvals_random = np.linalg.eigvalsh(H_mat)
E0 = min(eigvals_random)


# ## Constructing a variational circuit
# 
# The main task in VQE consists in designing a 'good' ansatz, i.e a family of variational states that can come close enough to the exact ground state. These states are constructed using quantum circuits with parameters (like rotation angles) that will be optimized to minimize the energy.

# In[ ]:


from qat.lang.AQASM import Program, QRoutine, RY, CNOT, RX, Z, H, RZ

def make_ansatz():
    """Function that prepares an ansatz circuit
    
    Returns:
        Circuit: a parameterized circuit (i.e with some variables that are not set)
    """
    
    nparams = ... # number of parameters 
    
    prog = Program()
    reg = prog.qalloc(nqbits)
    # define variables using 'new_var' 
    theta = [prog.new_var(float, '\\theta_%s'%i)
             for i in range(nparams)]
    
    # for instance...: put a rotation on each qubit
    for ind in range(nqbits):
        RY(theta[ind])(reg[ind])
        
    # and so on
    # ...
    
    circ = prog.to_circ()
    return circ

circ = make_ansatz()
circ.display()


# 
# ## Creating a variational job and a variational stack
# 
# We now create a variational job from this circuit and observable.
# It is then submitted to a "variational stack" composed of a perfect QPU, ``LinAlg``, and a variational plugin, ``ScipyMinimizePlugin``. The latter handles parametric jobs. These are jobs whose circuit contains a parameter that is then going to be optimized, using classical optimizers, so as to minimize the value of the observable over the final state.
# 
# Below, we are going to test three different classical optimizers: COBYLA, Nelder-Mead, and BFGS.

# In[ ]:


from qat.qpus import get_default_qpu
from qat.plugins import ScipyMinimizePlugin #,SPSAMinimizePlugin

circ = make_ansatz()

job = circ.to_job(job_type="OBS",
                  observable=hamiltonian,
                  nbshots=100)
# note that this job contains a circuit whose variables are not set!
# it therefore requires a special plugin so that the QPU can handle it!

theta_0 = ... # initial value of the parameters

linalg_qpu = get_default_qpu()
method = "COBYLA"  # look at the doc of scipy.optimize.minimize for other options
optimizer_scipy = ScipyMinimizePlugin(method=method,
                                      tol=1e-6,
                                      options={"maxiter": 200},
                                      x0=theta_0)
# you could also use SPSAMinimizePlugin instead

# we now build a stack that can handle variational jobs
qpu = optimizer_scipy | linalg_qpu

# we submit the job
result = qpu.submit(job)

print("Minimum VQE energy (%s) = %s"%(method, result.value))


# What do you observe? It the VQE energy close to the exact energy ``E0``? If not, where does it come from: a poor ansatz? a poor optimizer? a poor QPU? not enough ``nbshots``? Try playing with these various knobs. Hint: entanglement helps!
# 
# 
# ## Plotting the results
# 
# We can also plot the value of the variational energy over the course of the classical optimization. For this, we can retrieve information about the variational job execution in the ``meta_data`` field of the result.

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
plt.plot(eval(result.meta_data['optimization_trace']),
         label=method)
    
plt.plot([E0 for _ in range(300)], '--k', lw=3, label="exact energy")
    
plt.grid()
plt.legend(loc="best")
plt.xlabel("Steps")
plt.ylabel("Energy");


# What can you observe in terms of convergence speed?

# Once you have found an ansatz and an optimizer that works with a perfect QPU, can we now test what happens with a noisy QPU?

# ## With a noisy QPU
# 
# We can now execute the circuit on a noisy QPU simulator. To achieve this, we are going to turn our perfect QPU ``qpu`` into a noisy one by simply combining it with a so-called plugin that adds "depolarizing noise" on the circuit:

# In[ ]:


from depolarizing_plugin import DepolarizingPluginVec
depol_plugin = DepolarizingPluginVec(prob_1qb=0.005, prob_2qb=0.05)
qpu = get_default_qpu()
noisy_qpu = depol_plugin | qpu


# Depolarizing noise, as any noise, is defined by its action on the density matrix $\rho$:
# 
# $$\mathcal{E}(\rho) = (1-p) \rho + \frac{p}{3} \left( X \rho X + Y \rho Y + Z \rho Z \right)$$
# 
# with $X$, $Y$, and $Z$ the three Pauli matrices, and $p$ the probability that a depolarization event happens. (This is for one-qubit depolarizing noise; a similar expression holds for two-qubit depolarizing noise, with $X$, $Y$ and $Z$ replaced by tensor products of Pauli matrices).
# 
# In other words, depolarizing noise leaves the qubit unchanged with probability $1-p$, and otherwise applies a $X$, $Y$ or $Z$ gate with probability $p/3$. $p$ is what is called ``prob_1qb`` above. 
# 
# This type of noise is not exactly the type of noise that is observed in actual quantum computers, but it is close enough and it allows to understand the main effects of noise.
# 
# Here, the plugin adds depoalizing noise after each gate to mimic an imperfect gate.
# 
# - Check that this noise preserves the trace of the density matrix
# - What are the corresponding Kraus operators?
# 
# We now execute the circuit on this noisy QPU.
# 
# What happens when you change the ``linalg_qpu`` object above to ``noisy_qpu``? What happens when you change the depolarizing probabilities?

# ### Bonus
# How could you test the robustness of VQE against systematic errors, without explicitly modifying the ``DepolarizingPlugin``?

# In[ ]: