import numpy as np


class Matrice(object):
    """
    Le système d'équation est construit sur la base d'une matrice A et un vecteur B modélisant le système, tels que
    A.X = B

    Chaque ligne de la matrice correspond à une équation, correspondant elle-même à un élément du graphe.
    Les équations sont dans le même ordre que les éléments du graphe.

    Chaque colonne correspond à une inconnue, également prise dans le même ordre que les éléments du graphe.
    Si l'élément est un Potentiel ou un noeud, l'inconnue est une tension.
    Si l'élément est une branche, l'inconnue est l'intensité parcourant la branche.
    """
    def __init__(self, graphe):
        self.vecteur = graphe.vecteur
        self.mat_A = None
        self.mat_B = None  # Equation matricielle AX = B
        self.n = len(self.vecteur)

    def init_mat(self):
        self.mat_A = np.zeros((self.n, self.n))
        self.mat_B = np.zeros(self.n)

        for ligne, element in enumerate(self.vecteur):  # Potentiel constant
            # assert ligne == self.vecteur.index(element)
            if element.type_potentiel == element.P_CST:
                # Potentiel constant : on ajoute une équation indiquant la valeur de la tension pour cette variable
                self.mat_A[ligne][ligne] = 1.0  # à noter que par construction, ligne == index
                self.mat_B[ligne] = element.u()

            elif element.type_potentiel == element.P_NOEUD:  # On applique la loi des noeuds
                # Noeud : on ajoute une équation sur la somme des intensités
                for (voisin, polarite, _) in element.voisins():
                    self.mat_A[ligne][self.vecteur.index(voisin)] = polarite
                self.mat_B[ligne] = 0.0

            elif element.type_potentiel == element.P_DLT:  # On applique la loi d'Ohm
                self.mat_A[ligne][ligne] = 1.0
                self.mat_A[ligne][self.vecteur.index(element.amont)] = -1.0 / element.r()
                self.mat_A[ligne][self.vecteur.index(element.aval)] = 1.0 / element.r()
                self.mat_B[ligne] = 0.0

            elif element.type_potentiel == element.P_COND:  # DU = q/C
                self.mat_A[ligne][self.vecteur.index(element.amont)] = 1.0
                self.mat_A[ligne][self.vecteur.index(element.aval)] = -1.0
                self.mat_B[ligne] = element.coef_mul_tension * element.q / element.capacite

            else:
                assert False, "Rien à faire là"

    def solve(self):
        """
        Résout le système d'équation et retourne le vecteur solution.
        """
        return np.linalg.solve(self.mat_A, self.mat_B)