import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt

# System parameters
m1, m2 = 1.0, 1.5  # Masses
k1, k2, k3 = 100.0, 150.0, 50.0  # Spring constants

# Mass and stiffness matrices
M = np.array([[m1, 0],
              [0, m2]])
K = np.array([[k1 + k2, -k2],
              [-k2, k2 + k3]])

# Eigenvalue problem for natural frequencies and mode shapes
eigenvalues, eigenvectors = eigh(K, M)
natural_frequencies = np.sqrt(eigenvalues)

print("Natural Frequencies:", natural_frequencies)
print("Mode Shapes (not normalized):\n", eigenvectors)

# Plot mode shapes with displacements along the springs, without normalizing the
# mode shapes
plt.figure(figsize=(10, 6))
for i, mode_shape in enumerate(eigenvectors.T):  # Transpose to iterate over columns (each mode shape)
    # Add zero displacement points for the fixed walls
    mode_shape_with_boundaries = np.concatenate(([0], mode_shape, [0]))
    x_positions = [0, 1, 2, 3]  # Positions of the walls and masses
    
    # Plot mode shape
    plt.plot(x_positions, mode_shape_with_boundaries, marker="o", linestyle="-", label=f"Mode {i+1}")
    
    # Annotate displacements for masses
    plt.text(1, mode_shape[0], f"{mode_shape[0]:.2f}", ha='center', va='bottom')
    plt.text(2, mode_shape[1], f"{mode_shape[1]:.2f}", ha='center', va='bottom')
    
    # Draw arrows only between the masses and between masses and the fixed points
    plt.arrow(1, mode_shape_with_boundaries[1], 1, mode_shape_with_boundaries[2] - mode_shape_with_boundaries[1],
              color='gray', head_width=0.05, length_includes_head=True)
    plt.arrow(2, mode_shape_with_boundaries[2], 1, mode_shape_with_boundaries[3] - mode_shape_with_boundaries[2],
              color='gray', head_width=0.05, length_includes_head=True)

plt.title("Mode Shapes with Boundary Conditions and Spring Displacements (Unnormalized)")
plt.xlabel("Position (Wall, Mass 1, Mass 2, Wall)")
plt.ylabel("Amplitude (Unnormalized)")
plt.xticks([0, 1, 2, 3], ["Fixed Wall", "Mass 1", "Mass 2", "Fixed Wall"])
plt.legend()
plt.grid()
plt.show()