import numpy as np
from IPython.display import Latex
from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector, schmidt_decomposition

ζ = Statevector([1/np.sqrt(2),0,0,1/np.sqrt(2)])
ζ.draw('latex',prefix='|\\zeta\\rangle = ')

ζ_sd = schmidt_decomposition(ζ,[0])
print(ζ_sd)

for i, (s,u,v) in enumerate(ζ_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {s}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

w = Statevector([0,1/np.sqrt(3),1/np.sqrt(3),0,1/np.sqrt(3),0,0,0])
w.draw('latex',prefix='|w\\rangle = ')

w_sd = schmidt_decomposition(w,[0,1])

for i, (s,u,v) in enumerate(w_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {s}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

Statevector(sum([suv[0]*np.kron(suv[1],suv[2]) for suv in w_sd]))

χ = Statevector(np.array([1,0,0,0,1,0,0,0,1])*1/np.sqrt(3),dims=(3,3))
χ.draw('latex', prefix='|\\chi\\rangle = ',max_size=10)

w_sd = schmidt_decomposition(χ,[0])
for i, (s,u,v) in enumerate(w_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {s}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

dims = (2,3,4)
k = Statevector(np.arange(np.prod(dims)),dims=dims)
k.draw('latex',max_size=24)

k_sd = schmidt_decomposition(k,[0,2])
for i, (s,u,v) in enumerate(k_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {s}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

Statevector(sum([suv[0]*np.kron(suv[1],suv[2]) for suv in k_sd]),dims=dims).draw('latex',max_size=24)

ξ = Statevector.from_label('++')
ξ.draw('latex',prefix='|\\xi\\rangle = ')

ξ_sd = schmidt_decomposition(ξ,[0])
for i, (s,u,v) in enumerate(ξ_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {np.round(s,6)}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

from schmidt_decomposition_local import schmidt_decomposition as sd_loc

ξ_sd = sd_loc(ξ,[0])
for i, (s,u,v) in enumerate(ξ_sd):
    print('---')
    display(Latex(f"$$\\lambda_{i} = {np.round(s,6)}$$"))
    display(u.draw('latex',prefix=f'|u_{i}\\rangle = '))
    display(v.draw('latex',prefix=f'|v_{i}\\rangle = '))

