import numpy as np
from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from ising_kitaev import initialize_chain, trotter, rotate_to_measurement_basis, add_measurement, initialize_coupler

zeeman_ferro = 0.01
zeeman_para = 10

initial_config = np.array([zeeman_para, zeeman_ferro, zeeman_ferro, zeeman_ferro])

qreg = QuantumRegister(5)
creg = ClassicalRegister(3)
qcirc = QuantumCircuit(qreg, creg)

initialize_chain(qcirc, qreg, initial_config, 'logical_one')
initialize_coupler(qcirc, qreg)
qcirc.draw()

trotter(qcirc, qreg, initial_config, 1.4, 0.1, 1)
qcirc.draw()

trotter(qcirc, qreg, initial_config, 0.0, 0.1, 1)

rotate_to_measurement_basis(qcirc, qreg, [1, 2, 3])
add_measurement(qcirc, qreg, creg, [1, 2, 3])
qcirc.draw()

from qiskit import Aer, execute

backend = Aer.get_backend('qasm_simulator')
job = execute(qcirc, backend, shots=2000)

job.status()

result = job.result()
print(result.get_counts())

