# coding=utf-8
from numpy import dot as mul

from precondition import *


"""
2 lab
=================

PCG

полиномиальное ускорение

фиксируется число шагов
MZ = R

трёхчленная реккурентная формула(через полиномы чебышёва и лежандра)

http://www.wikiwand.com/en/Symmetric_successive_overrelaxation ssor
https://github.com/pyamg/pyamg/tree/master/pyamg/relaxation тут и Чёбышев есть

"""

n = 5
epsilon = 0.1
import numpy as np
from numpy.random import randn as random_matrix
from random import randint
from numpy import linalg
import scipy


def jacobi_precondition(n=10):
    A = random_matrix(n, n)
    D = np.tril(np.triu(A, 0))

    return A, D
steps = 10


def pcg(A, S, b):
    S_t = np.transpose(S)
    # array: [k-1, k, k+1]

    x = [None, random_matrix(n, 1), None]
    r = [b - mul(A, x[1]), b - mul(A, x[1]), None]
    z = [mul(mul(S_t, S), r[1]), mul(mul(S_t, S), r[1]), None]
    betta = [None, 0, 0]
    p = [None, z[1], None]
    alpha = [None, None, None]


    for k in xrange(1, steps):
        z[1] = mul(mul(S_t, S), r[1])

        betta[2] = sum(z[1] * r[1])[0] / sum(z[0], r[0])[0]
        

        p[2] = z[1] + betta[2][0] * p[1] # TODO error
        alpha[2] = sum(z[1] * r[1])[0] / sum(mul(A, p[2]) * p[2])[0]

        x[2] = x[1] + alpha[2] * p[2]
        r[2] = r[1] - alpha[2] * mul(A, p[2])
        print sum(r[2] * r[2])[0]
        if sum(r[2] * r[2])[0] <= epsilon:
        	print 'k: {d}'.format(k)
        	return x[2]
        # тут массивы сдвигаем на 1 влево:)
        x[:-1], r[:-1], z[:-1], betta[:-1], p[:-1], alpha[:-1] = x[1:], r[1:], z[1:], betta[1:], p[1:], alpha[1:]



    # [:-1] [1:]

    """
    x_k = random_matrix(n, 1)
    r_k = b - mul(A, x_k)
    betta = 0
    k = 0
    z_k = mul(mul(S_t, S), r_k)
    while sum(r_k * r_k) >= epsilon:
        z_1, r_1 = z_k, r_k
        z_k = mul(mul(S_t, S), r_k)
        if k >= 1:
            betta = sum(z_k * r_k) / sum(z_1, r_1)
            p_k = z_k + betta * p_k
        else:
            p_k = z_1
        p_k = z_k + betta * p_k
        alpha = sum(z_k * r_k) / sum(mul(A, p_k), p_k)
        x_k = x_k + alpha * p_k
        r_k = r_k - alpha * mul(A, p_k)

        k += 1

    """

def main():
    #A, S = jacobi_precondition(n=n)
    A, S = ssor_precondition(w=1.0, n=n)
    b = random_matrix(n, 1)
    #print linalg.solve(A,b)
    print pcg(A, S, b)


if __name__ == '__main__':
    main()