from .fast_math import op_ket, bj, commute





def time_evolve_com_xh(phi_in, x, ham, w0, dt, bj_list=None, eps=10e-5):
    phi_0 = phi_in
    phi_0_x = 0
    phi_1 = op_ket(ham, phi_in)
    c_xh = commute(x, ham)

    if bj_list is None:
        phi_1_x = op_ket(c_xh, phi_in)
        phi_out = 2 * (-1j) * bj(m=0, x=w0 * dt) * phi_1

        m = 2
        bjm = bj(m, w0 * dt)
        while abs(bjm) > eps:
            phi_2 = 2 * op_ket(ham, phi_1) - phi_0
            phi_2_x = 2 * op_ket(c_xh, phi_1) + 2 * op_ket(ham, phi_1_x) - phi_0_x
            phi_out = phi_out + 2 * (-1j ** m) * bjm * phi_2_x
            phi_0 = phi_1
            phi_1 = phi_2
            phi_0_x = phi_1_x
            phi_1_x = phi_2_x
            m += 1
            bjm = bj(m, w0 * dt)
        return phi_out

    else:

        phi_1_x = op_ket(c_xh, phi_in)
        phi_out = 2 * (-1j) * bj_list[0] * phi_1

        m = 2
        while abs(bj_list[m]) > eps:
            phi_2 = 2 * op_ket(ham, phi_1) - phi_0
            phi_2_x = 2 * op_ket(c_xh, phi_1) + 2 * op_ket(ham, phi_1_x) - phi_0_x
            phi_out = phi_out + 2 * (-1j ** m) * bj(m, x=w0 * dt) * phi_2_x
            phi_0 = phi_1
            phi_1 = phi_2
            phi_0_x = phi_1_x
            phi_1_x = phi_2_x

            m += 1