import numpy as np


def generateY_0():
    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():
    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():
    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():
    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(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():
    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(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():
    i0 = 100;
    G_i0 = np.dot(np.transpose(-generateQii(i0 + 1) - np.dot(generateQiiPlus1(), G)), generateQiiMin1(i0 + 1))
    while np.linalg.norm(G_i0 - G) < e_g:
        i0 += 10
        G_i0 = np.dot(np.transpose(-generateQii(i + 1) - np.dot(generateQiiPlus1(), G)), generateQiiMin1(i + 1))
    return i0;


def generateG_i(G_iPlus1, i):
    G_i = np.dot(np.linalg.inv(-generateQii(i + 1) - np.dot(generateQiiPlus1(), G_iPlus1)), generateQiiMin1(i + 1))
    return G_i;


def generateQ(i, Gi):
    return generateQii(i) + np.dot(generateQiiPlus1(), Gi)


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


def solveP():
    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) -> bool:
    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

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


if __name__ == '__main__':
    print('Hello! Please, enter W')
    # W = int(input())
    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)])

    print('Please, enter matrix elements')
    D_0 = np.matrix([[-8, 1], [1, -11]])
    D_1 = np.matrix([[2, 5], [4, 6]])

    # 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())
    M = 2
    print('Enter b')
    b = np.array([float(i) for i in input().split()])
    # 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 = 0.000000001
    e_f = 0.000001
    alpha = 100000000

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

    Y_0 = generateY_0()
    Y_1 = generateY_1()
    Y_2 = generateY_2()
    G = solveEquation()

    Q00 = generateQii(0)
    Q01 = generateQiiPlus1()
    Q10 = generateQiiMin1(1)

    I = find_i0()
    print("i0 = ", I)

    i = I
    listG = [G]
    while i >= 0:
        listG = [generateG_i(listG[0], i)] + listG
        i -= 1

    listQ = []
    i = I
    while i >= 0:
        listQ = [generateQ(i, listG[i])] + listQ
        i -= 1

    listF = [np.identity((W + 1) * (M + 1))]
    i = 1
    while i < I + 1:
        listF.append(generateF_i(i, listQ[i]))
        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(i, generateQ(i, G)))
            i += 1
    p0 = solveP()
    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
    print("Average count of requests on orbit", probSum)