import numpy as np
import matplotlib.pyplot as plt
# %matplotlib notebook

from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
from qiskit.opflow import StateFn, PauliSumOp

def generate_2qubit_qnn(label):
    circuit = QuantumCircuit(2)
    circuit.ry(np.pi / 4.0, 0)
    circuit.ry(np.pi / 4.0, 1)

    circuit.barrier()
    for i in range(2):
        if label[i] == 'x':
            circuit.rx(Parameter('theta'+str(i)), i)
        elif label[i] == 'y':
            circuit.ry(Parameter('theta'+str(i)), i)
        elif label[i] == 'z':
            circuit.rz(Parameter('theta'+str(i)), i) 
    circuit.barrier()

    # Add CZ ladder.
    circuit.cz(0, 1)

    return circuit

# # RealAmplitudesの場合
# from qiskit.circuit.library import RealAmplitudes

# qc_test = RealAmplitudes(num_qubits=2, reps=1,
#                         entanglement='linear').decompose()

qc_test = generate_2qubit_qnn('xy')
qc_test.draw('mpl')

# 期待値の計算
from qiskit import Aer
from qiskit.utils import QuantumInstance
from qiskit.opflow import CircuitSampler
from qiskit.opflow.expectations import PauliExpectation

hamiltonian = PauliSumOp.from_list([('ZZ', 1.0)])

expectation = StateFn(hamiltonian, is_measurement=True) @ StateFn(qc_test)
aerpauli_basis = PauliExpectation().convert(expectation)

quantum_instance = QuantumInstance(Aer.get_backend('qasm_simulator'),
                                  shots = 32768)
sampler = CircuitSampler(quantum_instance)

# 期待値の評価
def calculate_exp_val(params):
  value_dict = dict(zip(qc_test.parameters, params))
  result = sampler.convert(aerpauli_basis, params=value_dict).eval()
  return np.real(result)

# コスト関数の可視化
X = np.linspace(0, 2*np.pi, 30)
Y = np.linspace(0, 2*np.pi, 30)
Z = np.array([[calculate_exp_val([x, y]) for x in X] for y in Y]).reshape(len(X), len(Y))

xx, yy = np.meshgrid(X, Y)
fig = plt.figure(figsize=plt.figaspect(1))
ax = fig.add_subplot(projection="3d")
ax.plot_surface(xx, yy, Z, cmap=plt.cm.jet)

plt.show()

from qiskit.opflow import Gradient

gradient = Gradient().convert(expectation)
gradient_in_pauli_basis = PauliExpectation().convert(gradient)

# 勾配の評価
def calculate_gradient(params):
  value_dict = dict(zip(qc_test.parameters, params))
  result = sampler.convert(gradient_in_pauli_basis, params=value_dict).eval()
  return np.real(result)

# initial_point = 2*np.pi*np.random.random(qc_test.num_parameters)

# アニメーション用初期値
initial_point = np.array([np.pi, 2.0])

from qiskit.algorithms.optimizers import GradientDescent

gd_loss= []
xx_op= []
yy_op= []
def gd_callback(nfevs, x, fx, stepsize):
  gd_loss.append(fx)
  xx_op.append(x[0])
  yy_op.append(x[1])

gd = GradientDescent(maxiter=100,learning_rate=0.05, callback=gd_callback)

x_opt, fx_opt, nfevs = (
  gd.optimize(initial_point.size,
                calculate_exp_val,
                gradient_function=calculate_gradient,
                initial_point=initial_point)
)

from qiskit.algorithms.optimizers import SPSA

spsa_loss= []
xx_op_spsa= []
yy_op_spsa= []

# コールバック処理の定義
def spsa_callback(nfevs, x, fx, stepsize, accepted):
  spsa_loss.append(fx)
  xx_op_spsa.append(x[0])
  yy_op_spsa.append(x[1])

# インスタンス化
spsa = SPSA(maxiter=100,learning_rate=0.05,perturbation=0.05, callback=spsa_callback)

# 最適化処理の実行
x_opt_spsa, fx_opt_spsa, nfevs_spsa = (
  spsa.optimize(initial_point.size,
                calculate_exp_val,
                initial_point=initial_point)
)

# from qiskit.opflow import NaturalGradient

# ng = NaturalGradient().convert(expectation)
# ng_in_pauli_basis = PauliExpectation().convert(ng)

# def calculate_natural_gradient(params):
#   value_dict = dict(zip(qc_test.parameters, params))
#   result = sampler.convert(ng_in_pauli_basis, params=value_dict).eval()
#   return np.real(result)

# ngd_loss= []
# xx_op_ng= []
# yy_op_ng= []
# def ngd_callback(nfevs, x, fx, stepsize):
#   ngd_loss.append(fx)
#   xx_op_ng.append(x[0])
#   yy_op_ng.append(x[1])

# qng = GradientDescent(maxiter=50,learning_rate=0.1, callback=ngd_callback)

# x_opt_ng, fx_opt_ng, nfevs_ng = (
#   qng.optimize(initial_point.size,
#                 calculate_exp_val,
#                 gradient_function=calculate_natural_gradient,
#                 initial_point=initial_point)
# )

# from qiskit.algorithms.optimizers import QNSPSA

# qnspsa_loss = []
# xx_op_qnspsa= []
# yy_op_qnspsa= []
# def qnspsa_callback(nfev, x, fx, stepsize, accepted):
#     qnspsa_loss.append(fx)
#     xx_op_qnspsa.append(x[0])
#     yy_op_qnspsa.append(x[1])

# fidelity = QNSPSA.get_fidelity(qc_test,
#                                quantum_instance,
#                                expectation=PauliExpectation())
# qnspsa = QNSPSA(fidelity, maxiter=50, learning_rate=0.1,
#                                        perturbation=0.1,
#                                        callback=qnspsa_callback)

# x_opt_qnspsa, fx_opt_qnspsa, nfevs = qnspsa.optimize(initial_point.size, 
#                                        calculate_exp_val, 
#                                        initial_point=initial_point)

from matplotlib.animation import FuncAnimation

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.plot_surface(xx, yy, Z, cmap=plt.cm.jet, alpha=0.6)
gd, = ax.plot([], [], [], '.-k', label='Gradient Descent')
spsa, = ax.plot([], [], [], '.-r', label='SPSA')
# ngd, = ax.plot([], [], [], '.-b', label='Quantum Natural Gradient')
# spsa_ngd, = ax.plot([], [], [], '.-y', label='QN-SPSA')

ax.set_xlabel('theta_1')
ax.set_ylabel('theta_2')
ax.set_zlabel('loss')

def update(f):
    gd.set_data(xx_op[:f], yy_op[:f])
    gd.set_3d_properties(gd_loss[:f])
    spsa.set_data(xx_op_spsa[:f], yy_op_spsa[:f])
    spsa.set_3d_properties(spsa_loss[:f])
#     ngd.set_data(xx_op_ng[:f], yy_op_ng[:f])
#     ngd.set_3d_properties(ngd_loss[:f])
#     spsa_ngd.set_data(xx_op_qnspsa[:f], yy_op_qnspsa[:f])
#     spsa_ngd.set_3d_properties(qnspsa_loss[:f])
    
anim = FuncAnimation(fig, update, interval=200)

plt.legend(loc='upper left')
plt.show()

# 結果のプロット
plt.figure(figsize=(7, 3))
plt.plot(gd_loss, 'C0', label='vanilla gradient descent')
plt.plot(spsa_loss, 'C0', ls='--', label='SPSA')
plt.ylabel('loss')
plt.xlabel('iterations')
plt.legend()
plt.show()

anim.save("spsa.gif")

import qiskit.tools.jupyter
%qiskit_version_table