from sympy.utilities.pytest import XFAIL

from sympy.physics.quantum.qapply import qapply
from sympy.physics.quantum.qubit import Qubit
from sympy.physics.quantum.shor import (
    CMod, continued_fraction, getr, arr
)


@XFAIL
def test_CMod():
    assert qapply(CMod(4, 2, 2)*Qubit(0,0,1,0,0,0,0,0)) ==\
    Qubit(0,0,1,0,0,0,0,0)
    assert qapply(CMod(5, 5, 7)*Qubit(0,0,1,0,0,0,0,0,0,0)) ==\
    Qubit(0,0,1,0,0,0,0,0,1,0)
    assert qapply(CMod(3, 2, 3)*Qubit(0,1,0,0,0,0)) ==\
    Qubit(0,1,0,0,0,1)

def test_continued_frac():
    assert continued_fraction(3245, 10000) == [0,3,12,4,13]
    assert continued_fraction(1932, 2568) == [0, 1, 3, 26, 2]
    assert continued_fraction(6589, 2569) == [2, 1, 1, 3, 2, 1, 3, 1, 23]
    assert getr(513, 1024, 10) == 2
    assert getr(169, 1024, 11) == 6
    assert getr(314, 4096, 16) == 13
    assert arr(5,3) == [1,0,1]
    assert arr(4,3) == [1,0,0]
    assert arr(8,5) == [0,1,0,0,0]

