import math
import numpy as np


class Component:

    def init_attr(self):
        self.E = 30000
        self.B = 1000
        self.H = 150
        self.vertical_q = 10
        self.length = math.sqrt(
            (self.end_point[1] - self.start_point[1]) ** 2
            + (self.end_point[0] - self.start_point[0]) ** 2
        )
        if self.start_point[0] == self.end_point[0]:
            self.alpha = 0
        else:
            self.alpha = math.atan(
                (self.end_point[1] - self.start_point[1])
                / (self.end_point[0] - self.start_point[0])
            )

    def __init__(self, p1, p2) -> None:
        self.start_point = p1
        self.end_point = p2
        self.section_kind = "Rectangle"
        self.init_attr()

    def set_comp_attr(self, e, b, h):
        self.E = e
        self.B = b
        self.H = h

    def set_vertical_q(self, q):
        self.vertical_q = q

    def create_k(self):
        self.calculate_section_attr()
        K = np.zeros((6, 6))

        l = self.length
        i = self.E * self.I / l
        B = self.E * self.area / l
        c_x = math.cos(self.alpha)
        c_y = math.sin(self.alpha)
        a1 = B * c_x * c_x + 12 * i / l / l * c_y * c_y
        a2 = (B - 12 * i / l / l) * c_x * c_y
        a3 = B * c_y * c_y + 12 * i / l / l * c_x * c_x
        a4 = -6 * i / l * c_y
        a5 = 6 * i / l * c_x
        a6 = 4 * i
        K[0, 0] = K[3, 3] = a1
        K[0, 3] = K[3, 0] = -a1
        K[1, 0] = K[4, 3] = K[0, 1] = K[3, 4] = a2
        K[1, 3] = K[0, 4] = K[3, 1] = K[4, 0] = -a2
        K[1, 1] = K[4, 4] = a3
        K[1, 4] = K[4, 1] = -a3
        K[0, 2] = K[0, 5] = K[2, 0] = K[5, 0] = a4
        K[2, 3] = K[3, 5] = K[3, 2] = K[5, 3] = -a4
        K[1, 2] = K[1, 5] = K[2, 1] = K[5, 1] = a5
        K[2, 4] = K[4, 5] = K[4, 2] = K[5, 4] = -a5
        K[2, 2] = K[5, 5] = a6
        K[2, 5] = K[5, 2] = a6 / 2
        return K

    def create_f(self):
        F = np.zeros((6, 1))
        # 这里只针对竖向荷载
        length = abs(self.end_point[0] - self.start_point[0])

        F[1, 0] = -self.vertical_q * length / 2
        F[2, 0] = -self.vertical_q * length**2 / 12
        F[4, 0] = -self.vertical_q * length / 2
        F[5, 0] = self.vertical_q * length**2 / 12
        return F

    def calculate_section_attr(self):
        if self.section_kind == "Rectangle":
            self.I = self.B * self.H**3 / 12
            self.area = self.B * self.H

    def set_f(self, f_1_x, f_1_y, f_2_x, f_2_y):
        self.f_1_x = f_1_x / 1e3
        self.f_1_y = f_1_y / 1e3
        self.f_2_x = f_2_x / 1e3
        self.f_2_y = f_2_y / 1e3

    def set_m(self, m_1, m_2):
        self.m_1 = m_1 / 1e6
        self.m_2 = m_2 / 1e6

    @property
    def v1(self):
        return self.f_1_y * np.cos(self.alpha) - self.f_1_x * np.sin(self.alpha)

    @property
    def v2(self):
        return -self.f_2_y * np.cos(self.alpha) + self.f_2_x * np.sin(self.alpha)
