# define Hamiltonian
from qiskit.opflow import PauliSumOp
H = PauliSumOp.from_list([('XYII', 1), ('IYZI', 2), ('IIZX', 3), ('XIII', 4), ('IYII', 5)])

print(H)

from qiskit.circuit.library import EfficientSU2

qc = EfficientSU2(num_qubits=H.num_qubits, reps=1)

qc_with_meas = qc.measure_all(inplace=False)
qc_with_meas.decompose().draw(fold=120)

# define a set of (random) parameter values for the ansatz circuit
import numpy as np
theta = np.random.rand(qc.num_parameters)

# use the Sampler to sample from this circuit
from qiskit.primitives import Sampler

sampler = Sampler([qc_with_meas])
s_result = sampler([qc_with_meas], [theta])
print(s_result)

from qiskit.primitives import Estimator

estimator = Estimator([qc], [H])
e_result = estimator([qc], [H], [theta])
print(e_result)

from qiskit.algorithms import NumPyEigensolver

exact_solver = NumPyEigensolver(k=2) # find the first two '2' eigenvalues
exact_result = exact_solver.compute_eigenvalues(H)
print(exact_result.eigenvalues)

# define objective as expectation value of Hamiltonian with circuit ansatz

objective = lambda x: estimator([qc], [H], [x]).values[0]

# instantiate optimizer

from qiskit.algorithms.optimizers import SPSA
optimizer = SPSA(maxiter=500)

# define initial values for our circuit parameters
x0 = np.random.rand(qc.num_parameters)

# minimize the objective function
result = optimizer.minimize(objective, x0=x0)

# store ground state parameters for later
ground_state_params = result.x

# print the resulting ground state energy
print(result.fun)

from qiskit.algorithms import MinimumEigensolver, VQEResult

class CustomVQE(MinimumEigensolver):
    
    def __init__(self, estimator, circuit, optimizer, callback=None):
        self._estimator = estimator
        self._circuit = circuit
        self._optimizer = optimizer
        self._callback = callback
        
    def compute_minimum_eigenvalue(self, operator, aux_operators=None):
                
        # define objective
        def objective(x):
            e_job = self._estimator([self._circuit], [operator], [x])
            value = e_job.values[0]
            if self._callback:
                self._callback(value)
            return value
            
        # run optimization
        x0 = np.random.rand(self._circuit.num_parameters)
        res = self._optimizer.minimize(objective, x0=x0)
        
        # populate results
        result = VQEResult()
        result.cost_function_evals = res.nfev
        result.eigenvalue = res.fun
        result.optimal_parameters = res.x
        return result

# run the custom VQE function
custom_vqe = CustomVQE(estimator, qc, optimizer)
result = custom_vqe.compute_minimum_eigenvalue(H)
print(result)

from qiskit.circuit import ParameterVector

class Fidelity():
    
    def __init__(self, sampler):
        """ Instantiate fidelity estimator. """
        #self._sampler = sampler
        
    def run(self, circuit_1, circuit_2, param_values_1, param_values_2):
        """ Estimate fidelity between two states defined by given parameter values."""
        
        # prepare circuit
        circuit_1_ = circuit_1.remove_final_measurements(inplace=False)        
        circuit_2_ = circuit_2.remove_final_measurements(inplace=False)
        x = ParameterVector("x", circuit_1_.num_parameters)
        y = ParameterVector("y", circuit_2_.num_parameters)
        circuit = circuit_1_.assign_parameters(x).compose(circuit_2_.inverse().assign_parameters(y))
        circuit.measure_all()
        
        # run circuit and get probabilities
        param_values = np.append(param_values_1, param_values_2)
        sampler = Sampler([circuit])
        s_job = sampler([circuit], [param_values])
        probabilities = s_job.quasi_dists[0]

        # estimate fidelity: |<0|U_1^dag U_2|0>|^2 clipped to [0, 1]
        fidelity = np.minimum(1.0, np.maximum(0.0, probabilities.get(0, 0.0)))
        return fidelity

# instantiate fidelity
fidelity = Fidelity(sampler)

# interpolation between parameters
ts = np.linspace(0, 1)
fidelities = [fidelity.run(qc, qc, ground_state_params, ground_state_params + t*theta) for t in ts]

import matplotlib.pyplot as plt
plt.figure(figsize=(8, 4))
plt.plot(ts, fidelities)
plt.xlabel('t', fontsize=20)
plt.ylabel('fidelity', fontsize=20)
plt.show()

# set penalty weight for overlap term
penalty = 25

# define objective for VQD
vqd_objective = lambda x: objective(x) + penalty*fidelity.run(qc, qc, ground_state_params, x)

# run optimization to get first excited state
result = optimizer.minimize(vqd_objective, ground_state_params)
print(result)

# determine energy
e_job = estimator([qc], [H], [result.x])
print(e_job)

fid = fidelity.run(qc, qc, result.x, ground_state_params);
print(fid)

from qiskit.algorithms import Eigensolver, EigensolverResult

class VQD(Eigensolver):
    
    def __init__(self, estimator, fidelity, circuit, optimizer, penalty=25, callback=None):
        
        # self._estimator = estimator
        self._fidelity = fidelity
        self._circuit = circuit
        self._optimizer = optimizer
        self._penalty = penalty
        self._callback = callback
        
    def compute_eigenvalues(self, operator, aux_operators=None):
                
        # compute ground state
        estimator = Estimator([self._circuit], [operator])
        energy = lambda x: estimator([self._circuit], [operator], [x]).values[0]
        x0 = np.random.rand(self._circuit.num_parameters)
        res_1 = self._optimizer.minimize(energy, x0)

        # compute first excited states
        overlap = lambda x: self._fidelity.run(self._circuit, self._circuit, res_1.x, x)
        def objective(x):
            value = energy(x) + self._penalty * overlap(x)
            if self._callback:
                self._callback(value)
            return value
        res_2 = self._optimizer.minimize(objective, res_1.x)
            
        # populate results
        result = EigensolverResult()
        result.eigenvalues = [estimator([self._circuit], [operator], [res_1.x]),
                              estimator([self._circuit], [operator], [res_2.x])]
        return result        

vqd = VQD(estimator, fidelity, qc, optimizer=SPSA(maxiter=1000))
result = vqd.compute_eigenvalues(H)
print(result)

from qiskit_ibm_runtime import (Session, Options, QiskitRuntimeService,
                                Sampler as RuntimeSampler, 
                                Estimator as RuntimeEstimator)

service = QiskitRuntimeService(channel='ibm_quantum')

from qiskit.algorithms.optimizers import COBYLA
optimizer = COBYLA(maxiter=400)

with Session(service=service, backend='ibmq_qasm_simulator') as session:

    # prepare primitives
    rt_estimator = RuntimeEstimator(session=session)
    rt_sampler = RuntimeSampler(session=session)
    rt_fidelity = Fidelity(rt_sampler)

    # set up algorithm    
    rt_vqd = VQD(rt_estimator, rt_fidelity, qc, optimizer)
    
    # run algorithm
    result = rt_vqd.compute_eigenvalues(H)
    print(result)    

import qiskit_ibm_runtime
qiskit_ibm_runtime.version.get_version_info()

from qiskit.tools.jupyter import *

%qiskit_version_table
%qiskit_copyright