import numpy as np
from grid import Grid
from universal_element import UniversalElement
import shape_f
from data import Data
from jacobian import Jacobian
from matrixH import MatrixH
from matrixC import MatrixC
from matrixHBC import MatrixHBC


class FEM:
    def __init__(self, height, width, nodes_height, nodes_width):
        self.g = Grid(height, width, nodes_height, nodes_width)
        self.u_e = UniversalElement()
        self.d = Data()
        self.j = Jacobian(self.g, self.u_e, self.d)
        self.h = MatrixH(self.g, self.u_e, self.d, self.j)
        self.c = MatrixC(self.g, self.u_e, self.d, self.j)
        self.hbc = MatrixHBC(self.g, self.u_e, self.d, self.j, self.h)


        self.H_g = np.zeros((self.g.n_n, self.g.n_n))
        self.C_g = np.zeros((self.g.n_n, self.g.n_n))
        self.H_Cdt = np.zeros((self.g.n_n, self.g.n_n))
        self.calculate_h_g_c_g()
        self.calculate_h_cdt_g()

        self.p_for_side_pc = np.zeros((self.d.number_of_sides, self.d.number_of_pc_for_bc, self.u_e.num_of_shape_f,
                                       self.d.vector_columns))
        self.p_for_side = np.zeros((self.d.number_of_sides, self.u_e.num_of_shape_f, self.d.vector_columns))
        self.p_for_elements = np.zeros((self.g.n_e, self.u_e.num_of_shape_f, self.d.vector_columns))
        self.calculate_p_for_side_pc()
        self.calculate_p_for_sides()
        self.calculate_p_for_elements()

        self.P_g = np.zeros((self.g.n_n, 1))
        self.calculate_p_g()

        self.T_1 = np.zeros((self.g.n_n, 1))
        self.T_1.fill(100)
        self.P_Cdt_T0 = np.zeros((self.g.n_n, 1))
        self.calculate_p_cdt_t0()

        self.min_max_temperatures = np.zeros((self.d.number_of_steps, 1, 2))
        self.simulate()








    # AGREGACJA macierzy h i macierzy c (czyli dodawanie do odpowiednich komórek macierzy globalnej odpowiednich komó©ek macierzy lokalnych - po to mamy zapisane indeksy nodeów w elemetach) - powstaje macierz H i c dla całej siatki (o rozmairze ilość nodeów*ilość nodeów)
    def calculate_h_g_c_g(self):
        for i in range(self.g.n_e):
            i1 = 0
            for j in self.g.elements[i].nodes_id:
                i2 = 0
                for k in self.g.elements[i].nodes_id:
                    self.H_g[j][k] += self.hbc.h_plus_hbc[i][i1][i2]
                    self.C_g[j][k] += self.c.c_for_el[i][i1][i2]

                    i2 += 1
                i1 += 1

    # implementaja wzorku co jest w pdfach (to się nie zmienia, jest stale w symulacji)
    def calculate_h_cdt_g(self):
        self.H_Cdt = self.H_g + self.C_g/self.d.dT


    # wektor p liczy się jak hbc, więc opisałem powyżej (w sumie to to samo tyle że z minusem XD)
    # TODO: shorten the code using nested loops
    def calculate_p_for_side_pc(self):
        # POINT:   side / point / coordinate
        # MATRIX: side / point / shape_f
        # 1:
        self.p_for_side_pc[0][0][0] = -shape_f.n1(self.u_e.int_points_for_Hbc[0][0][0],
                                                  self.u_e.int_points_for_Hbc[0][0][1])
        self.p_for_side_pc[0][0][1] = -shape_f.n2(self.u_e.int_points_for_Hbc[0][0][0],
                                                  self.u_e.int_points_for_Hbc[0][0][1])

        self.p_for_side_pc[0][1][0] = -shape_f.n1(self.u_e.int_points_for_Hbc[0][1][0],
                                                  self.u_e.int_points_for_Hbc[0][1][1])
        self.p_for_side_pc[0][1][1] = -shape_f.n2(self.u_e.int_points_for_Hbc[0][1][0],
                                                  self.u_e.int_points_for_Hbc[0][1][1])
        # 2:
        self.p_for_side_pc[1][0][1] = -shape_f.n2(self.u_e.int_points_for_Hbc[1][0][0],
                                                  self.u_e.int_points_for_Hbc[1][0][1])
        self.p_for_side_pc[1][0][2] = -shape_f.n3(self.u_e.int_points_for_Hbc[1][0][0],
                                                  self.u_e.int_points_for_Hbc[1][0][1])

        self.p_for_side_pc[1][1][1] = -shape_f.n2(self.u_e.int_points_for_Hbc[1][1][0],
                                                  self.u_e.int_points_for_Hbc[1][1][1])
        self.p_for_side_pc[1][1][2] = -shape_f.n3(self.u_e.int_points_for_Hbc[1][1][0],
                                                  self.u_e.int_points_for_Hbc[1][1][1])
        # 3:
        self.p_for_side_pc[2][0][2] = -shape_f.n3(self.u_e.int_points_for_Hbc[2][0][0],
                                                  self.u_e.int_points_for_Hbc[2][0][1])
        self.p_for_side_pc[2][0][3] = -shape_f.n4(self.u_e.int_points_for_Hbc[2][0][0],
                                                  self.u_e.int_points_for_Hbc[2][0][1])

        self.p_for_side_pc[2][1][2] = -shape_f.n3(self.u_e.int_points_for_Hbc[2][1][0],
                                                  self.u_e.int_points_for_Hbc[2][1][1])
        self.p_for_side_pc[2][1][3] = -shape_f.n4(self.u_e.int_points_for_Hbc[2][1][0],
                                                  self.u_e.int_points_for_Hbc[2][1][1])
        # 4:
        self.p_for_side_pc[3][0][0] = -shape_f.n1(self.u_e.int_points_for_Hbc[3][0][0],
                                                  self.u_e.int_points_for_Hbc[3][0][1])
        self.p_for_side_pc[3][0][3] = -shape_f.n4(self.u_e.int_points_for_Hbc[3][0][0],
                                                  self.u_e.int_points_for_Hbc[3][0][1])

        self.p_for_side_pc[3][1][0] = -shape_f.n1(self.u_e.int_points_for_Hbc[3][1][0],
                                                  self.u_e.int_points_for_Hbc[3][1][1])
        self.p_for_side_pc[3][1][3] = -shape_f.n4(self.u_e.int_points_for_Hbc[3][1][0],
                                                  self.u_e.int_points_for_Hbc[3][1][1])

    # to liczę jak calculate_hbc_4x4 (opisane wyżej, tylko tu inne stałe są)
    def calculate_p_for_sides(self):
        for i in range(0, 4, 2):
            self.p_for_side[i] = (self.hbc.hbc_4x1[i][0]*self.d.weights_hbc[0] +
                                  self.hbc.hbc_4x1[i][1]*self.d.weights_hbc[1])* \
                                 self.d.alfa*self.d.t_ot*self.g.dx/2
            self.p_for_side[i + 1] = (self.hbc.hbc_4x1[i + 1][0]*self.d.weights_hbc[0] +
                                      self.hbc.hbc_4x1[i + 1][1]*self.d.weights_hbc[1])* \
                                     self.d.alfa*self.d.t_ot*self.g.dy/2

    # liczę dla każdego elementu wektor p czyli dodaję wektory p dla ścian do elementów co są na brzegch(a te co w środku mają wektor p = 0)
    def calculate_p_for_elements(self):
        j = self.g.n_e - 1
        for i in range(self.g.n_h - 1):
            self.p_for_elements[i] += self.p_for_side[3]
            self.p_for_elements[j] += self.p_for_side[1]
            j -= 1
        j = self.g.n_h - 2
        i = 0
        while i < self.g.n_e - 1:
            self.p_for_elements[i] += self.p_for_side[0]
            self.p_for_elements[j] += self.p_for_side[2]
            i += self.g.n_h - 1
            j += self.g.n_h - 1

    # agregacja wektora P (podobnie jak agregacja h i c)
    def calculate_p_g(self):
        for i in range(self.g.n_e):
            i1 = 0
            for j in self.g.elements[i].nodes_id:
                self.P_g[j] += self.p_for_elements[i][i1]
                i1 += 1

    # to się zmienia co iterację (ale na początku trza raz policzyć to sb odrębną funkcję napisałem xD)
    # T1 to jest wektor temperatur początkowych dla nodeów (w pdf to jest T0)
    def calculate_p_cdt_t0(self):
        self.P_Cdt_T0 = self.P_g + (self.C_g/self.d.dT).dot(self.T_1)

    # tu się liczy wektor T1 - czyli wektor temperatur po czasie dT, i go liczę dla każdej iteraci i tylko on się zmienia
    # wzorki masz wypisane w pdf co jest test case
    # a min max wyciąga min i max wartość z wektora p
    def simulate(self):
        for i in range(10):
            self.T_1 = np.linalg.inv(self.H_Cdt).dot((self.C_g/self.d.dT).dot(self.T_1) + self.P_g)
            self.min_max_temperatures[i] = [self.T_1.min(), self.T_1.max()]


f = FEM(0.1, 0.1, 4, 4)
print(f.min_max_temperatures)