#!/usr/bin/env python3

import numpy as np
from itertools import count


def example_array():
    return np.array([[ 3, 2, 1],
                     [ 2, 3, 2],
                     [ 1, 2, 3]], dtype=np.float64)

def inverse_power_method(A, sigma, eps, max_iter=None):
    """Determine an eigenvector of matrix A with an eigenvalue closest to sigma.

    A       -- an N by N square matrix
    sigma   -- parameter to select eigenvalue
    eps     -- accuracy
    """
    A = np.asarray(A)
    N = len(A)
    B = np.linalg.inv( A - sigma*np.identity(N))
    b_0 = np.ones(N)


    for i in count():
        b_1 = B @ b_0

        b_1 = b_1 / np.linalg.norm(b_1)
        
        error = np.sqrt(np.sum( (np.abs(b_0) - np.abs(b_1))**2 ))

        if error < eps:
            return b_1, i

        # Running out of iterations
        if max_iter is not None and max_iter >= i:
            raise RuntimeError("Not converging in {} steps".format(i))

        # Setup for next loop
        b_0 = b_1


def test_inverse_power_method():
    A = example_array() 
    sigma = 1
    eps = 1e-4
    max_iter = 1e5
    eigvec, n = inverse_power_method(A, sigma, eps, None)

    _, npeigvecs = np.linalg.eig(A)
   
    rtol = 1e-2
    found = False
    for npeigvec in npeigvecs.T:
        if np.allclose(npeigvec, eigvec, rtol):
            found = True

    assert found == True


def tridiagonalize(A):
    """ Compute the Tridiagonal of the matrix A. """
    return householder_tridiagonalize(A)

def householder_tridiagonalize(A):
    """ Compute the Tridiagonal of the matrix A using the Householder Method. """

    A = np.asarray(A)

    N = len(A)

    for k in range(N-1):
        q = np.sqrt(np.sum( A[k+1:,k]**2 ))
        alpha = -q * np.sign(A[k+1,k])
        r = np.sqrt( (alpha**2 - A[k+1,k]*alpha)/2 )

        v = np.zeros(N)
        v[k+1] = (A[k+1, k] - alpha)
        v[k+2:] = A[k+2:, k]
        v[k + 1:] /= 2 * r


        P = np.identity(N) - 2 * np.outer(v,v)

        A = P @ A @ P

    return A



def test_tridiagonalize():
    """ Compare the householder tridiagonalize function with Numpy's Hessenberg. """
    from scipy.linalg import hessenberg

    A = example_array() 
    T = tridiagonalize(A)
    H = hessenberg(A)

    rtol = 1e-3
    print(T)
    print(H)
    print(np.isclose(H, T, rtol))

    print("They are almost the same except for a minus sign at (2,3) and (3,2)")


def QREig(T, eps, max_iter=None):
    """Compute Eigenvalues using QR decomposition."""
    for i in count():
        Q, R = np.linalg.qr(T)

        T = R @ Q

        # We are done if the sub-diagonal is smaller than eps
        d_1 = np.diag(T, -1)
        if np.linalg.norm(d_1) < eps:
            return np.diagonal(T), i

        # Running out of iterations
        if max_iter is not None and max_iter >= i:
            raise RuntimeError("Not converging in {} steps".format(i))

def test_QREig():
    A = example_array()

    T = tridiagonalize(A)
    eps = 1e-6

    eigvals, _ = QREig(T, eps)
    npeigvals, npeigvecs = np.linalg.eig(T)

    assert np.allclose(eigvals, npeigvals, eps), "The computated eigenvalues do not match Numpy's eigenvalues"


def eig(A, QReps, PMeps, sigma=1e-2):
    """Determine eigenvectors with a combination of QR decomposition and inverse Power Method."""

    T = tridiagonalize(A)
    eigvals, _ = QREig(T, QReps)
    eigvecs = np.zeros((len(eigvals),len(A)), dtype=np.float64)

    sigma = (2*np.random.rand() -1 )*sigma
    for i, eigval in enumerate(eigvals):
        eigvecs[i], _ = inverse_power_method(T, eigval + sigma, PMeps)

    return eigvals, eigvecs 


def test_eig(debug=True):
    A = example_array()

    eps = 1e-5
    rtol = 1e-1
    eigvals, eigvecs = eig(A, eps, eps)

    if debug:
        for i, eigval in enumerate(eigvals):
            print("Eigval {}: ".format(eigval), eigvecs[i])

    npeigvals, npeigvecs = np.linalg.eig(A)
    if debug:
        for i, npeigval in enumerate(npeigvals):
            print("Numpy Eigval {}: ".format(npeigval), npeigvecs[i])

    assert np.allclose(eigvals, npeigvals, eps), "The computated eigenvalues do not match Numpy's eigenvalues"

    #assert np.allclose(eigvecs, npeigvecs, rtol),  "The computated eigenvectors do not match Numpy's eigenvectors"

    found = np.zeros(len(eigvecs))
    for i, eigvec in enumerate(eigvecs):
        for npeigvec in npeigvecs:
            if np.allclose(eigvec, npeigvec, rtol):
                found[i] = 1
                break

    print("Currently only the first eigenvector is the same from Numpy and my implementation. The others have a minus sign.")
    print("Comparable eigenvectors: ", found)

if __name__ == "__main__":
    #test_inverse_power_method()
    #test_tridiagonalize()
    #test_QREig()
    test_eig()