from .fast_math import op_ket, bj


def time_evolve_h(phi_in, ham, w0, dt, bj_list=None, eps=10e-5):
    phi_0 = phi_in
    phi_1 = op_ket(ham, phi_in)

    if bj_list is None:
        phi_out = bj(m=0, x=w0 * dt) * phi_0 + 2 * (-1j) * bj(m=1, 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_out = phi_out + 2 * (-1j ** m) * bj(m, x=w0 * dt) * phi_2
            phi_0 = phi_1
            phi_1 = phi_2
            m += 1
            bjm = bj(m, w0 * dt)
        return phi_out

    else:

        phi_out = bj_list[0] * phi_0 + 2 * (-1j) * bj_list[1] * phi_1
        m = 2
        while abs(bj_list[m]) > eps:
            phi_2 = 2 * op_ket(ham, phi_1) - phi_0
            phi_out = phi_out + 2 * (-1j ** m) * bj_list[m] * phi_2
            phi_0 = phi_1
            phi_1 = phi_2
            m += 1
        return phi_out