import numpy as np


def eigen_ub(m, sup=True):
    n_row, n_col = m.shape
    if sup:
        evs, _ = np.linalg.eig(m)
        return max(evs)
    else:
        return np.sqrt(n_row * n_col) * np.abs(m).max()


def eigen_ub_power_method(m, conv_thresh=0.001, max_iter=10000, print_interval=1):
    print('estimating max eigenvalue')
    x = np.random.randn(m.shape[0])
    x = x / np.sqrt(x @ x)
    prev_ev = 0
    for ii in range(max_iter):
        y = m.dot(x)
        x = y / np.sqrt(y.dot(y))
        ev = x.dot(m.dot(x))

        abs_ev = abs(ev)
        diff = abs(ev - prev_ev)
        if ii % print_interval == 0:
            print(f'iteraton: {ii} eigenvalue: {ev} diff: {diff}')
        if diff < conv_thresh * abs_ev:
            print(f'power method converged at iteration {ii}')
            return ev
        else:
            prev_ev = ev


def dot_inv_gd(m, v, lr=0.01, order=100, t_ini=None):
    n = len(m)
    assert m.shape == (n, n)

    if t_ini is None:
        ret = np.random.randn(*v.shape)
    else:
        ret = t_ini
    for _ in range(1, order + 1):
        ret = ret - lr * (m.dot(ret) - v)
        
    return ret


def dot_generalized_inv_gd(m, v, lr=0.01, order=100, t_ini=None):
    n = len(m)
    assert m.shape == (n, n)

    theta = 2 * np.pi / order
    cos = np.cos(theta)
    sin = np.sin(theta)

    if t_ini is None:
        # re = np.random.randn(len(v))
        re = np.zeros_like(v)
        im = np.zeros_like(v)
    else:
        re, im = t_ini
    
    for _ in range(1, order + 1):
        prev_re = re
        prev_im = im

        re = cos * prev_re - sin * prev_im
        re = re - lr * (m.dot(re) - v)
        im = cos * prev_im + sin * prev_re
        im = im - lr * m.dot(im)
              
    return re, im


def repeat_dot_inv_gd(m, v, algo, order=100, max_repeat=100, acc=0.001,
                      print_interval=1):
    ub = eigen_ub_power_method(m)
    assert ub > 0, 'estimated max eigenvalue is not positive'
    lr = 1 / ub
    print(f'lr: {lr}')
    ii = 0
    ret = v
    _ret = None
    for jj in range(1, max_repeat + 1):
        ii = jj
        prev_ret = ret
        _ret = algo(m, v, lr=lr, order=order, t_ini=_ret)
        if isinstance(_ret, tuple):
            ret = _ret[0]
        else:
            ret = _ret
        
        diff = (np.abs(ret - prev_ret)).max()
        norm = np.linalg.norm(ret)
        if ii % print_interval == 0:
            print(f'iteraton: {ii*order} norm: {norm} diff: {diff}')
        if norm == np.inf:
            lr /= 10.
            ret = v
            _ret = None
            print('matrix norm diverged.')
            print(f'lr: {lr}')
        elif diff < acc * norm:
            break
    
    print(f'Taylor series ended at iteration {ii*order}')
    return ret
