import numpy as np
import sympy as sym
from sympy import exp # Symbolic exponentiation.
from scipy.linalg import expm # Numerical exponentiation.
from sympy.physics.quantum import TensorProduct, OuterProduct

from qiskit import Aer, QuantumRegister, QuantumCircuit, assemble, execute
from qiskit.quantum_info import Statevector
from qiskit.circuit import Parameter
from qiskit.opflow import I, X, Y, Z

import pylatexenc # Required to use 'MatplotlibDrawer'.

'''1-qubit gates SU(2):'''
sigma0 = np.identity(2) # Matrix of the identity gate.
sigma1 = np.array([[0,1],[1,0]], dtype=(np.float32)) # Matrix of the Pauli-X gate that performs a Pi radian rotation around the x-axis.
sigma2 = np.array([[0,-1j],[1j,0]], dtype=(np.complex64)) # Matrix of the Pauli-Y gate = iXZ.
sigma3 = np.array([[1,0],[0,-1]], dtype=(np.float32)) # Matrix of the Pauli-Z gate = P(pi).
had = (1/np.sqrt(2))*np.array([[1,1],[1,-1]]) # Matrix of the Hadamard gate that performs a Pi radian rotation around an axis between the x and z axes.
phase = np.array([[1,0],[0,1j]]) #  Matrix of the Phase gate S = P(pi/2) = square-root of Pauli-Z.

'''2-qubit gates SU(4):'''
XX = np.kron(sigma1,sigma1) # Matrix of the XX gate.
YY = np.kron(sigma2,sigma2) # Matrix of the YY gate.

'''Array dimensions:'''
sigma0.shape, sigma1.shape, sigma2.shape, sigma3.shape, had.shape, phase.shape, XX.shape, YY.shape

# Circuit for the ZZ(t) gate:

'''
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
'''

t = Parameter('t')

qc = QuantumCircuit(2, name='ZZ')
qc.cnot(0,1) # CX^{01}.
qc.rz(2 * t, 1) # R_z(2t).
qc.cnot(0,1) # CX^{01}.
qc.draw(output='mpl') # '!pip install pylatexenc' library is required to use 'MatplotlibDrawer'.

# Circuit for the XX(t) gate:

'''
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
'''

t = Parameter('t')

qc = QuantumCircuit(2, name='XX')
qc.h([0,1])
qc.cx(0, 1)
qc.rz(2 * t, 1)
qc.cx(0, 1)
qc.h([0,1])
qc.draw(output='mpl')

# Equivalent circuit for the XX(t) gate:

qc = QuantumCircuit(2, name='XX')
qc.ry(np.pi/2,[0,1])
qc.cnot(0,1)
qc.rz(2 * t, 1)
qc.cnot(0,1)
qc.ry(-np.pi/2,[0,1])

qc.draw(output='mpl')

# Circuit for the YY(t) gate:

'''
from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
'''

t = Parameter('t')

qr = QuantumRegister(2)
qc = QuantumCircuit(qr, name='YY')
qc.sdg([0,1])
qc.h([0,1])
qc.cx(0, 1)
qc.rz(2 * t, 1)
qc.cx(0, 1)
qc.h([0,1])
qc.s([0,1])
qc.draw(output='mpl')

# Equivalent circuit for the YY(t) gate:

qr = QuantumRegister(2)
qc = QuantumCircuit(qr, name='YY')
qc.rx(np.pi/2,[0,1])
qc.cnot(0,1)
qc.rz(2 * t, 1)
qc.cnot(0,1)
qc.rx(-np.pi/2,[0,1])

qc.draw(output='mpl')

t = Parameter('t')
qr = QuantumRegister(2)

ZZ_qc = QuantumCircuit(qr, name='ZZ')
ZZ_qc.cnot(0,1)
ZZ_qc.rz(2 * t, 1)
ZZ_qc.cnot(0,1)
ZZg = ZZ_qc.to_instruction()

XX_qc = QuantumCircuit(qr, name='XX')
XX_qc.h([0,1])
XX_qc.append(ZZg, [0,1])
XX_qc.h([0,1])
XXg = XX_qc.to_instruction()

YY_qc = QuantumCircuit(qr, name='YY')
YY_qc.sdg([0,1])
YY_qc.h([0,1])
YY_qc.append(ZZg, [0,1])
YY_qc.h([0,1])
YY_qc.s([0,1])
YYg = YY_qc.to_instruction()

num_qubits = 3
Trot_qr = QuantumRegister(num_qubits)
Trot_qc = QuantumCircuit(Trot_qr, name='Trot')
for i in range(0, num_qubits - 1):
    Trot_qc.append(YYg, [Trot_qr[i], Trot_qr[i+1]])
    Trot_qc.append(XXg, [Trot_qr[i], Trot_qr[i+1]])
Trot_gate = Trot_qc.to_instruction()
Trot_qc.draw(output='mpl')

XXI = np.kron(XX, sigma0)
YYI = np.kron(YY, sigma0)

(expm(-1j*(XXI+YYI)).round(6) == (expm(-1j*XXI)@expm(-1j*YYI)).round(6)).all()

((expm(-1j*XXI)@expm(-1j*YYI)).round(6) == (expm(-1j*YYI)@expm(-1j*XXI)).round(6)).all()

'''Return should be >>> False.'''

YYI = np.kron(YY, sigma0)
IXX = np.kron(sigma0, XX)

(expm(-1j*(YYI+IXX)).round(6) == (expm(-1j*YYI)@expm(-1j*IXX)).round(6)).all()

from qiskit.opflow import I, X, Y, Z

H1 = (X^X^I)+(Y^Y^I)
(H1.exp_i().to_matrix().round(7) == ((X^X^I).exp_i()@(Y^Y^I).exp_i()).to_matrix().round(7)).all()

'''Return should be >>> False.'''

H2 = (Y^Y^I)+(I^X^X)
(H2.exp_i().to_matrix().round(3) == ((Y^Y^I).exp_i()@(I^X^X).exp_i()).to_matrix().round(3)).all()