from qiskit import QuantumCircuit
import numpy as np

def sim_z(t, qc, qubits):
    """
    Add gates to simulate a Pauli Z string exponential
   
    Parameters:
        t: float
            Time parameter to simulate for
        qc: QuantumCircuit
            Circuit to append gates to
        qubits: array
            Array indicating qubits indeces (in order) to 
            append the gates to
    """
    for i in range(len(qubits) - 1):
        qc.cx(qubits[i], qubits[i + 1])
    qc.rz(-2 * t, qubits[-1])
    for i in range(len(qubits) - 1, 0, -1):
        qc.cx(qubits[i - 1], qubits[i])

qc = QuantumCircuit(3)
sim_z(np.pi, qc, [0, 1, 2])
qc.draw()

def sim_pauli(arr, t, qc, qubits):
    """
    Append gates to simulate any Pauli string
    
    Parameters:
        arr: array
            Array encoding the Pauli string
        t: float
            Time parameter to simulate for
        qc: QuantumCircuit
            Circuit to append gates to
        qubits: array
            Array indicating qubits indeces (in order) to 
            append the gates to
    """
    new_arr = []
    new_qub = []
    for idx in range(len(arr)):
        if arr[idx] != 'I':
            new_arr.append(arr[idx])
            new_qub.append(qubits[idx])

    h_y = 1 / np.sqrt(2) * np.array([[1, -1j], [1j, -1]])
    for i in range(len(new_arr)):
        if new_arr[i] == 'X':
            qc.h(new_qub[i])
        elif new_arr[i] == 'Y':
            qc.unitary(h_y, [new_qub[i]], r'$H_y$')
            
    sim_z(t, qc, new_qub)
    
    for i in range(len(new_arr)):
        if new_arr[i] == 'X':
            qc.h(new_qub[i])
        elif new_arr[i] == 'Y':
            qc.unitary(h_y, [new_qub[i]], r'$H_y$')

def sim_ham(hamiltonian, t, qc, qubits, trotter=1):
    """
    Simulates Hamiltonian given as Pauli string
    
    Parameters:
        hamiltonian: dict
            Dictionary encoding the hamiltonian with each
            Pauli product as a key with the coefficient as value
        t: float
            Time parameter to simulate for
        qc: QuantumCircuit
            Circuit to append gates to
        qubits: array
            Array indicating qubits indeces (in order) to 
            append the gates to
        
    """
    temp    = QuantumCircuit(len(qubits))
    delta_t = t / trotter
    for pauli in hamiltonian:
        sim_pauli(pauli, hamiltonian[pauli] * delta_t, temp, range(len(qubits)))
        
    for i in range(trotter):
        qc.compose(temp, qubits, inplace=True)

qc = QuantumCircuit(3)
sim_ham({"XZY": 2, "ZXX": 5, "YXZ": 2}, 1 / (2 * np.pi), qc, [0, 1, 2], trotter=1) 
qc = qc.reverse_bits() # reverse because of Qiskit's endianness
qc.draw()

from qiskit import Aer, execute
from sympy import Matrix

qc = QuantumCircuit(3)
sim_ham({"XZY": 2, "ZXX": 5, "YXZ": 2}, 1 / (2 * np.pi), qc, [0, 1, 2], trotter=50)
qc = qc.reverse_bits()

backend = Aer.get_backend('statevector_simulator')
result  = execute(qc, backend).result()
vec     = result.get_statevector()
vec     = vec / vec[0] # global phase
qvec    = vec / np.linalg.norm(vec) # normalize
Matrix(np.round(qvec, 5))

import scipy

# Start with |0> state
start    = np.zeros(2 ** 3)
start[0] = 1

# Get the matrix corresponding to some Pauli product
# This function can be optimized, but it is left as it
# is for clarity purposes
def get_pauli(string):
    init   = string[0]
    string = string[1:]
    if init == 'X':
        out = np.array([[0, 1], [1, 0]])
    elif init == 'Z':
        out = np.array([[1, 0], [0, -1]])
    elif init == 'Y':
        out = np.array([[0, -1j], [1j, 0]])
    else:
        out = np.eye(2)
    
    for p in string:
        if p == 'X':
            out = np.kron(out, np.array([[0, 1], [1, 0]]))
        elif p == 'Z':
            out = np.kron(out, np.array([[1, 0], [0, -1]]))
        elif p == 'Y':
            out = np.kron(out, np.array([[0, -1j], [1j, 0]]))
        else:
            out = np.kron(out, np.eye(2))

    return out

# Hamiltonian is calculated from the Pauli decomposition
decomp = {"XZY": 2, "ZXX": 5, "YXZ": 2}
H      = np.zeros((2 ** 3, 2 ** 3))
H      = H.astype('complex128')
for pauli in decomp:
    H += decomp[pauli] * get_pauli(pauli)
    
# Hamiltonian is exponentiated and we multiply the starting
# vector by it to get our result
simul  = scipy.linalg.expm(1j * H * (1 / (2 * np.pi)))
vec    = simul @ start
vec    = vec / vec[0] # global phase
cvec   = vec / np.linalg.norm(vec) # normalize
Matrix(np.round(cvec, 5))

from qiskit.quantum_info import state_fidelity

state_fidelity(qvec, cvec)