import itertools
import numpy as np
from functools import partial, reduce
from qiskit import BasicAer, QuantumRegister, execute
from qiskit.quantum_info import Pauli
from qiskit_aqua import Operator, get_aer_backend
from qiskit_aqua.components.initial_states import Custom
from scipy.optimize import minimize
np.set_printoptions(precision=3, suppress=True)

def pauli_x(qubit, coeff):
    eye = np.eye((n_qubits))
    return Operator([[coeff, Pauli(np.zeros(n_qubits), eye[qubit])]])

n_qubits = 2

Hm = reduce(lambda x, y: x+y,
            [pauli_x(i, 1) for i in range(n_qubits)])
Hm.to_matrix()

def pauli_z(qubit, coeff):
    eye = np.eye((n_qubits))
    return Operator([[coeff, Pauli(eye[qubit], np.zeros(n_qubits))]])


def product_pauli_z(q1, q2, coeff):
    eye = np.eye((n_qubits))
    return Operator([[coeff, Pauli(eye[q1], np.zeros(n_qubits)) * Pauli(eye[q2], np.zeros(n_qubits))]])

J = np.array([[0,1],[0,0]])
Hc = reduce(lambda x,y:x+y,
        [product_pauli_z(i,j, -J[i,j])
         for i,j in itertools.product(range(n_qubits), repeat=2)])
Hc.to_matrix()

n_iter = 10 # number of iterations of the optimization procedure
p = 2
beta = np.random.uniform(0, np.pi*2, p)
gamma = np.random.uniform(0, np.pi*2, p)

init_state_vect = [1 for i in range(2**n_qubits)]
init_state = Custom(n_qubits, state_vector=init_state_vect)

qr = QuantumRegister(n_qubits)
circuit_init = init_state.construct_circuit('circuit', qr)

def evolve(hamiltonian, angle, quantum_registers):
    return hamiltonian.evolve(None, angle, 'circuit', 1,
                              quantum_registers=quantum_registers,
                              expansion_mode='suzuki',
                              expansion_order=3)

def create_circuit(qr, gamma, beta, p):
    circuit_evolv = reduce(lambda x,y: x+y, [evolve(Hm, beta[i], qr) + evolve(Hc, gamma[i], qr)
                                             for i in range(p)])
    circuit = circuit_init + circuit_evolv
    return circuit

def evaluate_circuit(gamma_beta, qr, p):
    n = len(gamma_beta)//2
    circuit = create_circuit(qr, gamma_beta[:n], gamma_beta[n:], p)
    return np.real(Hc.eval("matrix", circuit, get_aer_backend('statevector_simulator'))[0])
evaluate = partial(evaluate_circuit, qr=qr, p=p)

result = minimize(evaluate, np.concatenate([gamma, beta]), method='L-BFGS-B')
result

circuit = create_circuit(qr, result['x'][:p], result['x'][p:], p)

backend = BasicAer.get_backend('statevector_simulator')
job = execute(circuit, backend)
state = np.asarray(job.result().get_statevector(circuit))
print(np.absolute(state))
print(np.angle(state))

Z0 = pauli_z(0, 1)
Z1 = pauli_z(1, 1)

print(Z0.eval("matrix", circuit, get_aer_backend('statevector_simulator'))[0])
print(Z1.eval("matrix", circuit, get_aer_backend('statevector_simulator'))[0])