import numpy as np

import matplotlib.pyplot as plt

# Hàm tính ma trận độ cứng của các phần tử

def stiffness_matrix(E, A, L):
k = E * A / L
return np.array([[k, -k],
[-k, k]])

# Hàm tính ma trận chuyển vị

def displacement_matrix(nodes):
n = len(nodes)
D = np.zeros((n*2, n*2))
for i in range(n):
D[i*2, i*2] = 1
D[i*2+1, i*2+1] = 1
return D

# Hàm tính biến dạng và nội lực

def calculate_deformation_and_internal_force(material, section, nodes, elements,
boundary_conditions, loads):
E = material['E']
A = section['A']

n = len(nodes)
K = np.zeros((n*2, n*2))
D = displacement_matrix(nodes)

# Tính toán ma trận độ cứng tổng

for i, element in enumerate(elements):
node1, node2 = element
x1, y1 = nodes[node1]
x2, y2 = nodes[node2]
L = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
k = stiffness_matrix(E, A, L)

indices = [node1*2, node1*2+1, node2*2, node2*2+1]

for ii, i in enumerate(indices):

for jj, j in enumerate(indices):
K[i, j] += k[ii, jj]

# Áp dụng điều kiện biên

for node, displacement in boundary_conditions.items():
if 'x' in displacement:
K[node*2, :] = 0
K[:, node*2] = 0
K[node*2, node*2] = 1
if 'y' in displacement:
K[node*2+1, :] = 0
K[:, node*2+1] = 0
K[node*2+1, node*2+1] = 1

# Tính toán vector lực tải

F = np.zeros(n*2)
for node, load in loads.items():
F[node*2] = load.get('x', 0)
F[node*2+1] = load.get('y', 0)
 # Giải hệ phương trình
U = np.linalg.solve(K, F)

# Tính toán nội lực và chuyển vị

internal_forces = []
displacements = []
for i, element in enumerate(elements):
node1, node2 = element
x1, y1 = nodes[node1]
x2, y2 = nodes[node2]
L = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)

indices = [node1*2, node1*2+1, node2*2, node2*2+1]

u_e = U[indices]

# Chuyển vị và nội lực trong phần tử

u = np.array([[x1, y1], [x2, y2]]) + np.reshape(u_e, (2, 2)).T
internal_force = E * A / L * (u_e[2] - u_e[0])

displacements.append(u)
internal_forces.append(internal_force)

return displacements, internal_forces

# Ví dụ dữ liệu
material = {'E': 200e9} # Độ cứng của vật liệu (Pa)
section = {'A': 0.01} # Diện tích mặt cắt (m^2)
nodes = {0: [0, 0], 1: [2, 0], 2: [4, 0]} # Tọa độ nút
elements = [(0, 1), (1, 2)] # Các phần tử (liên kết nút)
boundary_conditions = {0: {'x': 0, 'y': 0}, 2: {'y': 0}} # Điều kiện biên tại nút
loads = {1: {'x': 10000}} # Tải tại nút (N)

# Tính toán
displacements, internal_forces = calculate_deformation_and_internal_force(material,
section, nodes, elements, boundary_conditions, loads)

# In kết quả
for i, (u, f) in enumerate(zip(displacements, internal_forces)):
print(f"Element {i+1}:")
print(f"Displacement:\n{u}")
print(f"Internal force: {f} N\n")

# Vẽ biểu đồ
for i, (u, f) in enumerate(zip(displacements, internal_forces)):
node1, node2 = elements[i]
x1, y1 = nodes[node1]
x2, y2 = nodes[node2]

plt.plot([x1, x2], [y1, y2], label=f'Element {i+1}')

plt.scatter([x1, x2], [y1, y2], color='red')

plt.text(x1, y1, f'{u[0]}', fontsize=12, va='bottom')
plt.text(x2, y2, f'{u[1]}', fontsize=12, va='bottom')

mid_x = (x1 + x2) / 2

mid_y = (y1 + y2) / 2
plt.text(mid_x, mid_y, f'{f} N', fontsize=12, ha='center', va='bottom')

plt.xlabel('X Coordinate (m)')

plt.ylabel('Y Coordinate (m)')
plt.legend()
plt.grid(True)
plt.show()