__author__ = 'chunk'

import numpy as np
import math


def gen_hilbert(n):
    return np.array([1.0 / (i + j + 1) for i in range(n) for j in range(n)]).reshape((n, -1))


def gen_b(n):
    return np.array([1.0 / i for i in range(1, n + 1)])


def norm(l):
    s = 0
    for item in l:
        if np.absolute(item) > s:
            s = np.absolute(item)
    return s


def jacobi(A, b, epsilon=0.0001, x=None):
    n = len(A)
    assert n == len(b)

    if x is None:
        x = np.zeros(n, dtype=np.float32)
    x_old = None

    D = np.diag(A)
    R = A - np.diagflat(D)

    # for i in range(N):
    iter = 0
    while x_old is None or norm(x - x_old) >= epsilon:
        iter += 1
        x_old = np.copy(x)
        x = (b - np.dot(R, x_old)) / D
        print iter, x
    return x


def SOR(A, b, omega, epsilon=0.0001, x=None):
    n = len(A)
    assert n == len(b)

    if x is None:
        x = np.zeros(n, dtype=np.float32)
    x_old = None
    iter = 0
    while x_old is None or norm(x - x_old) >= epsilon:
        iter += 1
        x_old = np.copy(x)
        for i in range(n):
            s = 0.0
            for j in range(n):
                if j != i:
                    s += A[i][j] * x[j]
            x[i] = (1 - omega) * x[i] + omega * (b[i] - s) / A[i][i]

            # print iter, x
    return iter, x


def test():
    # A = [[10, 3, 1], [2, -10, 3], [1, 3, 10]]
    # b = [14, -5, 14]
    # print SOR(A, b, 1.25)
    # print jacobi(A, b)

    n = 10
    A = gen_hilbert(n)
    b = gen_b(n)
    # print jacobi(A, b)
    # return
    print SOR(A, b, 1.25)

    for w in np.linspace(0.5, 1.5, num=11):
        iter, x = SOR(A, b, w)
        # print "w:%f\titer:%d\tx:%s\t" % (w, iter, x)
        # print "norm(delta):%s\n" % (norm(np.array([1.0] + [0.0] * (n - 1)) - x))
        print "w:%f\titer:%d\tnorm(delta):%s\n" % (
        w, iter, norm(np.array([1.0] + [0.0] * (n - 1)) - x))


if __name__ == '__main__':
    test()