import numpy as np


class truss:
    n: int
    s: int
    nodes: np.ndarray
    externalForces: list[list]
    beams: np.ndarray
    beamForces: np.ndarray

    # n: number nodes
    # s: number beams
    # externalForces: list of externalForces each force is list with [Fx, Fy, NodeIndex] so a single force in y direction on node 3 would look like: [[0,F,3]] (zero indexed!!)
    # nodeCoordinates: noarray of nodes x and y coordinates: [[x1,y1],[x2,y2]]
    # beams: ndarray of beams every beam contains node indices it connects. e.g. beam0 connnects nodes 0 and 1 -> beams= [[0,1]]
    def __init__(
        self,
        s: int,
        n: int,
        externalForces: list,
        nodeCoordinates: np.ndarray,
        beams: np.ndarray,
    ):
        self.n = n
        self.s = s
        self.externalForces = externalForces
        self.nodes = nodeCoordinates
        self.beams = beams

    def solve(self):
        # coeficient matrices A in x and y dir. with  #beams rows and #nodes cols
        # Ax,Ay is padded because numpy requires square coefficient matrix
        # Ax and Ay are concatenated efter lineaer system is constructed
        Ax = np.zeros((self.n, 2 * self.n))  # lhsx
        Ay = np.zeros((self.n, 2 * self.n))  # lhsy

        # rhs vectors all zero or external force
        bx = np.zeros(self.n)
        by = np.zeros((self.n))
        for ef in self.externalForces:
            bx[ef[2]] = ef[0]
            by[ef[2]] = ef[1]

        # construct lgs
        for ni in range(len(self.nodes)):  # for every node index ni
            for b in np.where(self.beams == ni)[
                0
            ]:  # for every beam connected to node[ni]
                (x1, y1) = self.nodes[self.beams[b][0]]
                (x2, y2) = self.nodes[self.beams[b][1]]

                # find deltaX and deltaY of beam to calculate force components
                vy, vx = 0, 0
                # if (
                # b < (2 * ni) and ni > 0
                # ):  # find force direction for beams on left side of node
                if self.beams[b][0] != ni:
                    vy = y1 - y2
                    vx = x1 - x2
                else:  # different force direction for nodes on right side of node
                    vy = y2 - y1
                    vx = x2 - x1

                lv = np.sqrt(vx**2 + vy**2)  # |v|

                ux = ((1 / lv) if lv != 0 else 0) * vx  # (v^)x unit vector X
                uy = ((1 / lv) if lv != 0 else 0) * vy  # (v^)y unit vector Y

                Ax[ni, b] = ux  # koeffizient stabkraft F[b] an node[ni] in x richtung
                Ay[ni, b] = uy  # koeffizient stabkraft F[b] an node[ni] in y richtung

        A = np.concatenate((Ax, Ay), axis=0)
        b = np.concatenate((bx, by), axis=0)

        self.beamForces = np.linalg.lstsq(A, b)[0]

    def print(self):
        [
            print(f"F{i+1}: {n:#.3g} kN")
            for i, n in enumerate(self.beamForces)
            if n != 0
        ]  # ignore padding collumns


### sample useage
## TM1 Aufgabe 164
if __name__ == "__main__":
    n = 15  # nr. nodes
    s = 2 * n - 3  # 27 beams(0-26) + F[27]+ Fa[28] + Fb[29]

    F = 28
    Fa = (((n - 3) / 2) * F) / 2
    Fb = Fa

    ## Node x and y coordinates
    nodes = np.array([[i, 0 if i % 2 == 0 else 2] for i in range(n)])

    # ExternalForces (x,y,nodeIndex)
    externalForces = [
        [0, Fa, 0],
        [0, -F, 2],
        [0, -F, 4],
        [0, -F, 6],
        [0, -F, 8],
        [0, -F, 10],
        [0, -F, 12],
        [0, Fb, 14],
    ]

    ## beams each row represents a beam and contains the node indices it connects to
    beams = np.array(
        [
            [
                x / 2 if x % 2 == 0 else (x - 1) / 2,
                x / 2 + 1 if x % 2 == 0 else (x - 1) / 2 + 2,
            ]
            for x in range(s)
        ],
        dtype=int,
    )

    t = truss(s, n, externalForces, nodes, beams)

    t.solve()
    t.print()