import numpy as np
import matplotlib.pyplot as plt

# W = 1
# b = np.array([1., 0.])
# M = 2
# D_0 = np.matrix([[-8, 1], [1, -11]])
# D_1 = np.matrix([[2, 5], [4, 6]])

def generateY_0(W, M, b):
    IB = np.kron(np.identity(W + 1), b)
    y_0 = np.matrix([[0 for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    y_0[0:W + 1, W + 1:(W + 1) * (M + 1)] = IB

    return y_0


def generateY_1(W, M, S, S_0, D_0):
    kron_sum = np.kron(D_0, np.identity(M)) + np.kron(np.identity(W + 1), S)

    diag = kron_sum.diagonal()
    R = np.diagflat(np.absolute(diag))
    y_21 = np.dot(np.linalg.inv(R), np.kron(np.identity(W + 1), S_0))
    tmp = np.dot(np.linalg.inv(R), kron_sum)
    y_22 = tmp + np.identity(tmp.shape[0])

    y_1 = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    y_1[W + 1: W + (W + 1) * (M + 1), 0: (W + 1)] = y_21
    y_1[W + 1:(W + 1) * (M + 1), W + 1:(W + 1) * (M + 1)] = y_22
    return y_1


def generateY_2(M, W, S, D_0, D_1):
    y_2 = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    kron_sum = np.kron(D_0, np.identity(M)) + np.kron(np.identity(W + 1), S)
    diag = kron_sum.diagonal()
    R = np.diagflat(np.absolute(diag))
    y_22 = np.dot(np.linalg.inv(R), np.kron(D_1, np.identity(M)))
    y_2[W + 1:(W + 1) * (M + 1), W + 1:(W + 1) * (M + 1)] = y_22
    return y_2


def solveEquation(W, M, Y_0, Y_1, Y_2, e_g):
    G_0 = np.identity((W + 1) * (M + 1))
    tmp1 = np.linalg.inv(np.identity((W + 1) * (M + 1)) - Y_1)
    tmp2 = Y_0 + np.dot(Y_2, np.dot(G_0, G_0))
    G_1 = np.dot(tmp1, tmp2)
    while np.linalg.norm(G_1 - G_0) > e_g:
        G_0 = G_1
        tmp1 = np.linalg.inv(np.identity((W + 1) * (M + 1)) - Y_1)
        tmp2 = Y_0 + np.dot(Y_2, np.dot(G_0, G_0))
        G_1 = np.dot(tmp1, tmp2)

    return G_1


def generateQii(W, M, S_0, D_0, D_1, b, alpha, S, i):
    q00 = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    q00[W + 1:(W + 1) * (M + 1), 0:W + 1] = np.kron(np.identity(W + 1), S_0)
    q00[0:(W + 1), 0:W + 1] = D_0 - i * alpha * np.identity(W + 1)
    q00[0:(W + 1), W + 1:(W + 1) * (M + 1)] = np.kron(D_1, b)
    kron_sum = np.kron(D_0, np.identity(M)) + np.kron(np.identity(W + 1), S)
    q00[W + 1:(W + 1) * (M + 1), W + 1:(W + 1) * (M + 1)] = kron_sum
    return q00


def generateQiiPlus1(W, M, D_1):
    q01 = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    q01[W + 1:(W + 1) * (M + 1), W + 1:(W + 1) * (M + 1)] = np.kron(D_1, np.identity(M));
    return q01


def generateQiiMin1(W, M, alpha, b, i):
    q10 = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    q10[0:(W + 1), W + 1:(W + 1) * (M + 1)] = np.kron(i * alpha * np.identity(W + 1), b)
    return q10


def find_i0(W, M, S_0, D_0, D_1, b, alpha, S, G, e_g):
    i0 = 10
    G_i0 = np.dot(np.transpose(-generateQii(W, M, S_0, D_0, D_1, b, alpha, S, i0 + 1) - np.dot(generateQiiPlus1(W, M, D_1), G)), generateQiiMin1(W, M, alpha, b, i0 + 1))
    while np.linalg.norm(G_i0 - G) < e_g:
        i0 += 10
        G_i0 = np.dot(np.transpose(-generateQii(W, M, S_0, D_0, D_1, b, alpha, S, i + 1) - np.dot(generateQiiPlus1(W, M, D_1), G)), generateQiiMin1(W, M, alpha, b, i + 1))
    return i0


def generateG_i(W, M, S_0, D_0, D_1, b, alpha, S, G_iPlus1, i):
    G_i = np.dot(np.linalg.inv(-generateQii(W, M, S_0, D_0, D_1, b, alpha, S, i + 1) - np.dot(generateQiiPlus1(W, M, D_1), G_iPlus1)), generateQiiMin1(W, M, alpha, b, i + 1))
    return G_i


def generateQ(W, M, S_0, D_0, D_1, b, alpha, S, i, Gi):
    return generateQii(W, M, S_0, D_0, D_1, b, alpha, S, i) + np.dot(generateQiiPlus1(W, M, D_1), Gi)


def generateF_i(W, M, D_1, i, Q, listF):
    return np.dot(np.dot(listF[i - 1], generateQiiPlus1(W, M, D_1)), np.linalg.inv(-Q))


def solveP(W, M, listQ, listF):
    Q00 = -listQ[0]
    F = np.matrix([[0. for __ in range((W + 1) * (M + 1))] for _ in range((W + 1) * (M + 1))])
    for f in listF:
        F += f
    F = np.dot(F, np.ones((M + 1) * (W + 1)))
    Q00[0:1, 0:(W + 1) * (M + 1)] = F
    zero = np.zeros((W + 1) * (M + 1))
    zero[0] = 1
    return np.linalg.solve(Q00, zero)


def criterion(D_0, D_1, S, beta, size):
    D_0 = np.array(D_0)
    D_1 = np.array(D_1)
    S = np.array(S)
    beta = np.array(beta)
    EPS = 10 ** (-5)
    nul = np.array([0] * size)
    eden = np.array([1] * size)
    # print(EPS)
    D = D_0 + D_1
    D = np.matrix.transpose(D)
    D[0] = [1] * size
    nul[0] = 1
    # print(D)
    Teta = np.linalg.solve(D, nul)
    print("Teta: ", Teta)
    lamb = np.dot(Teta, np.dot(D_1, eden))
    print("Lambda = ", lamb)
    # print("Inverse matrix S")
    # print(np.linalg.inv(S))
    # print(np.linalg.inv(S)*np.matrix.transpose(S))
    nym = np.dot(np.dot(beta, np.linalg.inv(S)), eden)
    nym = (-1) / nym
    print("Ny = ", nym)
    print("ro = ", lamb / nym)
    return (lamb / nym + EPS) < 1, lamb


def checkNegative(matrix):
    for i in range(len(matrix) - 1):
        if matrix[i, i] > 0:
            return False
    return True


def solution(i, alpha):
    print('Hello! Please, enter W')
    # W = int(input())
    M = 2
    W = 1
    D_0 = np.matrix([[0 for __ in range(W + 1)] for _ in range(W + 1)])
    D_1 = np.matrix([[0 for __ in range(W + 1)] for _ in range(W + 1)])
    D_0 = np.matrix([[-8, 1], [1, -11]])
    D_1 = i * np.matrix([[2, 5], [4, 6]])


    print('Please, enter matrix elements')

    # for i in range(W + 1):
    #     D_0[i] = [float(j) for j in input().split()]
    # for i in range(W + 1):
    #     D_1[i] = [float(j) for j in input().split()]

    print('Please, enter M')
    # M = int(input())
    print('Enter b')
    b = np.array([1, 0])
    # print('Enter S')
    # S = np.matrix([[0 for _ in range(M + 1)] for __ in range(M + 1)])

    S = np.matrix([[-30., 30.], [0., -30.]])
    E = np.ones(W + 1)
    S_0_1 = np.dot(-S, E)
    S_0 = np.transpose(S_0_1)

    # for i in range(M + 1):
    #     S[i] = [float(i) for i in input().split()]
    # e_g, e_f = [float(i) for i in input().split()]

    e_g = 10**-9
    e_f = 0.000001

    condition, lamb = criterion(D_0, D_1, S, b, W + 1)
    print(condition)

    Y_0 = generateY_0(W, M, b)
    Y_1 = generateY_1(W, M, S, S_0, D_0)
    Y_2 = generateY_2(M, W, S, D_0, D_1)
    G = solveEquation(W, M, Y_0, Y_1, Y_2, e_g)

    Q00 = generateQii(W, M, S_0, D_0, D_1, b, alpha, S, 0)
    Q01 = generateQiiPlus1(W, M, D_1)
    Q10 = generateQiiMin1(W, M, alpha, b, 1)

    I = find_i0(W, M, S_0, D_0, D_1, b, alpha, S, G, e_g)
    print("i0 = ", I)

    i = I
    listG = [G]
    while i >= 0:
        listG = [generateG_i(W, M, S_0, D_0, D_1, b, alpha, S, listG[0], i)] + listG
        i -= 1

    listQ = []
    i = I
    while i >= 0:
        listQ = [generateQ(W, M, S_0, D_0, D_1, b, alpha, S, i, listG[i])] + listQ
        i -= 1

    listF = [np.identity((W + 1) * (M + 1))]
    i = 1
    while i < I + 1:
        listF.append(generateF_i(W, M, D_1, i, listQ[i], listF))
        if np.linalg.norm(listF[len(listF) - 1], np.inf) < e_f:
            break
        i += 1

    if np.linalg.norm(listF[len(listF) - 1], np.inf) > e_f:
        while np.linalg.norm(listF[len(listF) - 1], np.inf) > e_f:
            listF.append(generateF_i(W, M, D_1, i, generateQ(W, M, S_0, D_0, D_1, b, alpha, S, i, G), listF))
            i += 1
    p0 = solveP(W, M, listQ, listF)
    p = []
    for f in listF:
        p.append(np.array(np.dot(p0, f)).flatten())

    i = 0
    prob = []
    for Pi in p:
        prob.append(np.dot(Pi, np.transpose(np.ones((W + 1) * (M + 1)))))
        print("probability that in orbit ", i, "applications: ", np.dot(Pi, np.transpose(np.ones((W + 1) * (M + 1)))))
        i += 1

    probSum = 0
    i = 1
    for pr in prob:
        probSum += i * pr
        i += 1
    return probSum, lamb


if __name__ == '__main__':
    Leng = [[], [], []]
    Lamb = [[], [], []]
    q = 0
    for alpha in [5, 10, 20]:
        L = list()
        La = list()
        koef = 1
        for i in range(1, 11):
            koef += 0.01
            length, lamb = solution(koef, alpha)
            L.append(length)
            La.append(lamb)
        Leng[q] = L
        Lamb[q] = La
        q += 1
    plt.ylabel(r'$L$')
    plt.xlabel(r'$\lambda$')
    print('Leng', Leng[0])
    print('lam', Lamb[0])
    print('Leng', Leng[0])
    print('lam', Lamb[1])
    plt.plot(Lamb[0], Leng[0], label=r'$\gamma=1$', linestyle='-')
    plt.plot(Lamb[1], Leng[1], label=r'$\gamma=5$', linestyle='--')
    plt.plot(Lamb[2], Leng[2], label=r'$\gamma=30$', linestyle='-.')
    plt.legend(loc='best')
    plt.show()