import numpy as np
from qiskit.quantum_info import state_fidelity

def schmidt_decomp_even(state):
    n   = int(np.log2(len(state)))
    mat = state.reshape(2 ** (n // 2), 2 ** (n // 2))
    
    u, s, v = np.linalg.svd(mat)
    
    a = [u[:, i].reshape(2 ** (n // 2), 1) for i in range(2 ** (n // 2))]
    b = [v[i, :].reshape(2 ** (n // 2), 1) for i in range(2 ** (n // 2))]    
    
    new_state = np.sum([np.kron(a[i], b[i]) * s[i] for i in range(2 ** (n // 2))], axis=0)    
    assert np.isclose(state_fidelity(state, new_state), 1), "Didnt work :("
    
    return a, b, s

n       = 2 # must be even
state   = np.random.rand(2 ** n)
state   = state / np.linalg.norm(state)
a, b, s = schmidt_decomp_even(state)

from qiskit import QuantumCircuit, transpile
from qiskit.quantum_info import Statevector, random_statevector
from sympy import Matrix

qc = QuantumCircuit(n)
qc.ry(2 * np.arccos(s[0]), 0)
qc.draw()

Matrix(np.round(Statevector(qc), 5))

qc.cx(0, 1)
qc.draw()

Matrix(np.round(Statevector(qc), 5))

theta = 2 * np.arccos(a[0][0][0])
phi   = np.real(np.log(a[0][1][0] / np.sin(np.arccos(a[0][0][0])), dtype='complex') / 1j)
lamb  = np.real(np.log(-a[1][0][0] / np.sin(np.arccos(a[0][0][0])), dtype='complex') / 1j)

qc.u(theta, phi, lamb, 0)
qc.draw()

Matrix(np.round(Statevector(qc), 5))

theta = 2 * np.arccos(b[0][0][0])
phi   = np.real(np.log(b[0][1][0] / np.sin(np.arccos(b[0][0][0])), dtype='complex') / 1j)
lamb  = np.real(np.log(-b[1][0][0] / np.sin(np.arccos(b[0][0][0])), dtype='complex') / 1j)

qc.u(theta, phi, lamb, 1)
qc.draw()

Matrix(np.round(Statevector(qc.reverse_bits()), 5))

Matrix(np.round(state, 5))

def two_qubit_state(state):
    qc = QuantumCircuit(2)
    
    # Normalize state and get Schmidt decomposition
    state   = state / np.linalg.norm(state)
    a, b, s = schmidt_decomp_even(state)

    # Phase 1
    qc.ry(2 * np.arccos(s[0]), 0)
    
    # Phase 2
    qc.cx(0, 1)
    
    # Phase 3
    qc.unitary(np.block([a[0].flatten().T, a[1].flatten().T]).reshape(2, 2).T, 0)
    
    # Phase 4
    qc.unitary(np.block([b[0].flatten().T, b[1].flatten().T]).reshape(2, 2).T, 1)
    
    return qc

state = random_statevector(4).data
qc    = two_qubit_state(state)
qc.decompose().draw()

circ = Statevector(qc.reverse_bits()).data
Matrix(np.round(circ, 5))

Matrix(np.round(state, 5))

for _ in range(100):
    state = random_statevector(4).data
    qc    = two_qubit_state(state)
    circ  = Statevector(qc.reverse_bits()).data
        
    assert np.allclose(circ, state)

