import numpy as np
import matplotlib.pyplot as plt
import PyQt5 as qt


def make_Y(n):
    Y = np.zeros((n, n))
    return Y

def connect_Y(x, y, Ys, Yp, Y):
    Y[x, y] = -Ys
    Y[y, x] = -Ys
    Y[x, x] += Ys + Yp
    Y[y, y] += Ys + Yp

def line_coor(n):
    tab = []
    i = 0
    k = 1
    while i < n:
        
        while k < n:
            tab += [[i, k]]
            k += 1
        i += 1
        k = i + 1

    return tab

def line_matrix(Y):
    n = len(Y)
    t = line_coor(n)
    lS = np.zeros((len(t), n))

    for i in range(len(t)):
        e = t[i]
        y = - Y[e[0], e[1]]
        lS[i, e[0]] = y
        lS[i, e[1]] = - y

    return lS

def delta_select(i, S):
    S = np.delete(S, (i), axis=0)
    S = np.delete(S, (i), axis=1)
    return S

def power_select(i, P):
    P = np.array(P[:i].tolist() + P[i+1:].tolist())
    return P

def complete_data(P, delta, i):
    ndelta = np.array(delta[:i].tolist() + [0] + delta[i:].tolist())
    nP = np.array(P[:i].tolist() + [-np.sum(P)] + P[i:].tolist())
    return ndelta, nP

def demo1():
    # Vn
    Vn = 2e5

    # Vecteur des puissances
    P = np.array([1000, -500, -250, -250])
    P = P * 1e6 # Passage en MW

    # Création de la matrice d'admitances (dimension n)
    Y = make_Y(4)
    connect_Y(2, 3, 0.1, 0, Y)
    connect_Y(1, 3, 0.15, 0, Y)
    connect_Y(2, 1, 0.05, 0, Y)
    connect_Y(2, 0, 0.05, 0, Y)
    connect_Y(3, 0, 0.05, 0, Y)
    print("Admittance matrix :", Y)

    # Mise en place du système linéaire à résoudre
    S = Y
    S = delta_select(3, S) # dim n-1, sélection de l'angle de transport de référence (delta_3)
    S *= Vn**2
    print("System matrix :", S)

    # Sélection des puissances (dimension n-1)
    P = power_select(3, P)
    print("Power input : ", P)

    # Résolution (dimension n-1)
    invS = np.linalg.inv(S)
    print("Inverse : ", invS)

    # Calcul des angles de transport (dimension n-1)
    delta = np.dot(invS, P)

    # Ajout de l'angle de transport d'origine et de la puissance associée (on repasse en dim n)
    ndelta, nP = complete_data(P, delta, 3)
    print("Power :", nP)
    print("Delta (deg) :", ndelta * 180 / 3.1415)

    # Calcul de la matrice de ligne
    lS = line_matrix(Y)
    print("Line matrix : ", lS)

    # Calcul des puissances de lignes
    line_power = Vn**2 * np.dot(lS, ndelta)
    lcoor = line_coor(len(ndelta))
    disp_line = []
    for i in range(len(line_power)):
        disp_line += [lcoor[i] + [line_power[i]]]
    print("Line power : ", disp_line)

def demo2():
    # Vn
    Vn = 4e5

    # Donnée de dimension
    Dim = 3
    NodeRef = 2

    # Vecteur des puissances
    P = np.array([500, 500, -1000])
    P = P * 1e6 # Passage en MW

    # Création de la matrice d'admitances (dimension n)
    Y = make_Y(Dim)
    connect_Y(0, 1, 25e-3, 0, Y)
    connect_Y(0, 2, 20e-3, 0, Y)
    connect_Y(1, 2, 50e-3, 0, Y)
    print("Admittance matrix : \n", Y)

    # Mise en place du système linéaire à résoudre
    S = Y # dim n
    S = delta_select(NodeRef, S) # dim n-1, sélection de l'angle de transport de référence (delta_3)
    S *= Vn**2
    print("System matrix : \n", S)

    # Sélection des puissances (dimension n-1)
    P = power_select(NodeRef, P)
    print("Power input : \n", P)

    # Résolution (dimension n-1)
    invS = np.linalg.inv(S)
    print("Inverse : \n", invS)

    # Calcul des angles de transport (dimension n-1)
    delta = np.dot(invS, P)

    # Ajout de l'angle de transport d'origine et de la puissance associée (on repasse en dim n)
    ndelta, nP = complete_data(P, delta, NodeRef)
    print("Power : \n", nP)
    print("Delta (deg) : \n", ndelta * 180 / 3.1415)

    # Calcul de la matrice de ligne
    lS = line_matrix(Y)
    print("Line matrix : \n", lS)

    # Calcul des puissances de lignes
    line_power = Vn**2 * np.dot(lS, ndelta)
    lcoor = line_coor(len(ndelta))
    disp_line = []
    for i in range(len(line_power)):
        disp_line += [lcoor[i] + [int(line_power[i] * 1e-6)]]
    print("Line power : \n", disp_line)

if __name__=="__main__":
    # Exemple du cours
    demo2()