#!/usr/bin/env python3

import numpy as np

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

def jacobi(A, b, eps):
    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)


    D = np.diag(A)
    L = -np.tril(A, -1)
    U = -np.triu(A, 1)

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

    # initially x_f = x_(i-1)
    # this changes when in the loop
    x_i = np.dot(D_inv, b)
    x_f = np.zeros(len(A))
    k = 1

    while diff(x_i, x_f) >= eps:
        k += 1

        # Save the previous solution vector as x_f
        x_f = x_i

        # Create new solution vector
        x_i = np.dot(np.dot(D_inv, ( L + U )), x_f ) + np.dot(D_inv, b)

    return x_i, k

def steepest_descent(A, b, eps):
    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)
   
    # initially x_f = x_(i-1)
    # this changes when in the loop
    x_f = np.zeros(len(A), dtype=np.float64)
    k = 1

    v_f = b
    t = np.dot(v_f, v_f) / np.dot(v_f, np.dot(A, v_f))
    x_i = x_f + t*v_f

    while diff(x_i, x_f) >= eps:
        k += 1

        # Pre calculate v_f and t
        v_f = b - np.dot(A, x_i)
        
        t = np.dot(v_f, v_f) / np.dot(v_f, np.dot(A, v_f))

        # Save the previous solution vector as x_f
        x_f = x_i

        # Create new solution vector
        x_i = x_f + t * v_f

    return x_i, k

def conjugate_gradient(A, b, eps):
    A = np.array(A, dtype=np.float64)
    b = np.array(b, dtype=np.float64)
   
    # initially x_f = x_(i-1)
    # this changes when in the loop
    x_f = np.zeros(len(A), dtype=np.float64)
    r_f = b - np.dot(A, x_f)
    v_f = r_f

    k = 1

    # Calculate first iteration
    t = np.dot(r_f, r_f) / np.dot(v_f, np.dot(A, v_f))
    x_i = x_f + t*v_f

    r_i = r_f - t * np.dot(A, v_f)
    s = np.dot(r_i, r_i) / np.dot(r_f, r_f)

    v_i = r_i - s*v_f

    # Set r and v vectors for next loop
    r_f = r_i
    v_f = v_i

    while diff(x_i, x_f) >= eps:
        k += 1
        
        t = np.dot(r_f, r_f) / np.dot(v_f, np.dot(A, v_f))
        # Save the previous solution vector as x_f
        x_f = x_i

        # Create new solution vector
        x_i = x_f + t*v_f

        # Calculate r and v vectors

        r_i = r_f - t * np.dot(A, v_f)
        s = np.dot(r_i, r_i) / np.dot(r_f, r_f)
        v_i = r_i - s*v_f
        
        # Save r and v vectors for next loop
        r_f = r_i
        v_f = v_i


    return x_i, k