"""
The LRM is based on the Lanczos algorithm for tridiagonalizing sparse Hermitian matrices.
The Lanczos algorithm is usually used to obtain extremal eigenvalues and the corresponding eigenstates,
 but it can also be used to calculate spectral property
"""

from .fast_math import op_ket, bra_ket, init_v0, init_vrand
import math


def compute_lanczos(ham, m_limit):
    b1 = 0
    phi_0 = init_v0(ham.shape[1])
    phi_1 = init_vrand(phi_0.shape)
    a_ = []
    b_ = [b1]
    for m in range(1, m_limit):
        psi_m = op_ket(ham, phi_1) - b_[m - 1] * phi_0
        a_m = bra_ket(psi_m, phi_1)
        a_.append(a_m)
        psi_m = psi_m - a_m * phi_1
        b_m1 = math.sqrt(bra_ket(psi_m, psi_m))
        b_.append(b_m1)

        phi_0 = phi_1
        phi_1 = psi_m / b_m1

    return [a_, b_]