#!/usr/bin/env python3

import numpy as np
from itertools import count as count

def diff(a, b):
    return np.amax(np.abs(a-b))

def jacobi(A, b, eps, max_iter = None):
    """ Use the Jacobi Method to solve a Linear System. """

    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)

    # Determine Diagonal and Upper and Lower matrices
    D = np.diag(A)
    L = -np.tril(A, -1)
    U = -np.triu(A, 1)

    D_inv = np.diagflat(np.reciprocal(D))

    x_0 = D_inv @ b
    for i in count():
        x_1 = D_inv @ ( L + U) @ x_0

        # Are we close enough?
        if diff(x_0, x_1) < eps:
            return x_1, i

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

        # Set values for next loop
        x_0 = x_1

def steepest_descent(A, b, eps, max_iter = None):
    """ Use Steepest Descent to solve a Linear System. """
    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)
   

    x_0 = np.zeros(len(A), dtype=np.float64)

    for i in count():
        Ax = A @ x_0
        v = b - Ax
        t = np.dot(v,v) / np.dot(v, A @ v )

        x_1 = x_0 + t*v

        # Are we close enough?
        if diff(x_0, x_1) < eps:
            return x_1, i

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

        # Set values for next loop
        x_0 = x_1

def conjugate_gradient(A, b, eps, max_iter = None):
    """ Use the Conjugate Gradient Method to solve a Linear System. """
    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)


    # Setup vectors
    x_0 = np.zeros(len(A), dtype=np.float64)
    r_0 = b - A @ x_0
    v = r_0.copy()

    for i in count():
        Ax = A @ x_0
        Av = A @ v
        
        r_0_square = np.dot(r_0, r_0)

        t = r_0_square / np.dot(v, Av )
        x_1 = x_0 + t*v

        r_1 = r_0 - t * Av

        s = np.dot(r_1, r_1) / r_0_square
        v = r_1 + s*v


        # Are we close enough?
        if diff(x_0, x_1) < eps:
            return x_1, i

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

        # Set values for next loop
        x_0 = x_1
        r_0 = r_1