import numpy as np
from thewalrus.random import random_symplectic
from thewalrus.quantum import is_symplectic
from strawberryfields.decompositions import bloch_messiah

r = np.array([1,1,1,1,2,2,0,0,0,0,0,0])
Sdiag= np.diag(np.concatenate([np.exp(r), np.exp(-r)]))
num_modes = len(r)
O1 = random_symplectic(num_modes, passive=True)
O2 = random_symplectic(num_modes, passive=True)
S = O1 @ Sdiag @ O2
assert is_symplectic(S)

# Do the Bloch-Messiah from SF
from strawberryfields.decompositions import bloch_messiah
O1, S, O2 = bloch_messiah(S)

assert is_symplectic(Sdiag)
assert is_symplectic(O1)
assert is_symplectic(O2)