import numpy as np
from matplotlib import pyplot as plt
from numpy.linalg import solve

# Constants
g = 9.81  # Acceleration due to gravity (m/s^2)

# Vehicle parameters
mass = 310.0             # kg
wd = 0.5                 # Front weight distribution (fraction)

# Applied acceleration inputs in G's (longitudinal, lateral, vertical)
# Replace these values with the desired inputs
long_acc_g = 0        # Longitudinal acceleration in G's (positive for acceleration, negative for braking)
lat_acc_g = 0      # Lateral acceleration in G's
vert_acc_g = 0        # Vertical acceleration in G's (e.g., due to bumps; usually small)

# Calculated accelerations in m/s^2
long_acc = long_acc_g * g  # Longitudinal acceleration (m/s^2)
lat_acc = lat_acc_g * g    # Lateral acceleration (m/s^2)
vert_acc = vert_acc_g * g  # Vertical acceleration (m/s^2)

# Geometry parameters (convert from mm to meters)
cm_m = 280 / 1000.0   # Center of mass height (m)
tw_m = 1200 / 1000.0  # Track width (m)
wb_m = 1550 / 1000.0  # Wheelbase (m)
wheel_center_z = 203.2 / 1000.0  # Wheel center height (m)

# Suspension attachment points - upright side (converted from mm to meters)
fb_upright = np.array([3.7628, -37.0299, 108.4733]) / 1000.0  # Fore Bottom
ab_upright = np.array([3.7628, -37.0299, 108.4733]) / 1000.0  # Aft Bottom
ft_upright = np.array([-3.7628, -49.3078, 297.9267]) / 1000.0  # Fore Top
at_upright = np.array([-3.7628, -49.3078, 297.9267]) / 1000.0  # Aft Top
d_upright = np.array([-4.5, -86.1697, 264.1451]) / 1000.0      # Damper
tr_upright = np.array([52.3129, -39.3953, 144.9733]) / 1000.0   # Tie Rod

# Suspension attachment points - chassis side (converted from mm to meters)
fb_chassis = np.array([208.68, -386.8966, 124.4613]) / 1000.0
ab_chassis = np.array([-150, -386.8966, 155.2616]) / 1000.0
ft_chassis = np.array([106.6266, -341.9893, 283.17]) / 1000.0
at_chassis = np.array([-125, -341.9893, 263.0367]) / 1000.0
d_chassis = np.array([-4.5, -270.4791, 95.2367]) / 1000.0
tr_chassis = np.array([81.6177, -417.12, 183.0733]) / 1000.0

# Contact Patch position (origin)
cP = np.array([0.0, 0.0, 0.0])  # m

# Applied forces at contact patch (N)
Fz_total = mass * (g + vert_acc)  # Total vertical load considering vertical acceleration
Fz_front_total = wd * Fz_total  # Total front axle vertical load (N)
Fz_front_per_wheel = Fz_front_total / 2  # Static vertical load per front wheel (N)

# Calculate weight transfer due to longitudinal and lateral acceleration
# Longitudinal weight transfer (ΔFz_long) (N)
wtLong = (mass * long_acc * cm_m) / wb_m

# Lateral weight transfer (ΔFz_lat) (N)
wtLat = (mass * lat_acc * cm_m) / tw_m

# Vertical load on the front outer wheel (N)
# For simplicity, assuming right turn (weight transfer to left wheel)
Fz_outer = Fz_front_per_wheel + (wtLat / 2) + (wtLong / 2)

# Applied vertical force at the contact patch (N)
Fz_wheel = Fz_outer

# Lateral and longitudinal forces at the wheel (N)
# Assuming friction coefficient (μ) approximated by lateral acceleration in G's
mu_lat = lat_acc_g
F_lat_wheel = Fz_wheel * mu_lat

mu_long = long_acc_g
F_long_wheel = Fz_wheel * mu_long

# Total applied force vector at the contact patch
fap = np.array([F_long_wheel, F_lat_wheel, Fz_wheel])  # Applied force at contact patch (N)
# fap = np.array([F_long_wheel, 2275, 1472])  # Applied force at contact patch (N)

# Compute position vector from contact patch to wheel center
wheel_center = np.array([0.0, 0.0, wheel_center_z])  # Wheel center position (m)
r_moment_arm = wheel_center - cP  # Position vector from contact patch to wheel center

# Compute moment due to applied force at contact patch about wheel center
mom = np.cross(r_moment_arm, fap)  # Moment due to applied force at contact patch about wheel center (N*m)

# Compute deltas and lengths for each suspension member
def compute_deltas_and_unit_vectors(upright, chassis):
    delta = chassis - upright
    length = np.linalg.norm(delta)
    unit_vector = delta / length
    return delta, unit_vector, length

# Fore Bottom Arm
deltaA, forceA, lA = compute_deltas_and_unit_vectors(fb_upright, fb_chassis)
# Aft Bottom Arm
deltaB, forceB, lB = compute_deltas_and_unit_vectors(ab_upright, ab_chassis)
# Fore Top Arm
deltaC, forceC, lC = compute_deltas_and_unit_vectors(ft_upright, ft_chassis)
# Aft Top Arm
deltaD, forceD, lD = compute_deltas_and_unit_vectors(at_upright, at_chassis)
# Damper
deltaE, forceE, lE = compute_deltas_and_unit_vectors(d_upright, d_chassis)
# Tie Rod
deltaF, forceF, lF = compute_deltas_and_unit_vectors(tr_upright, tr_chassis)

# Compute position vectors from wheel center to upright attachments
rA = fb_upright - wheel_center
rB = ab_upright - wheel_center
rC = ft_upright - wheel_center
rD = at_upright - wheel_center
rE = d_upright - wheel_center
rF = tr_upright - wheel_center

# Assemble force vectors and position vectors
forces_list = [forceA, forceB, forceC, forceD, forceE, forceF]
positions_list = [rA, rB, rC, rD, rE, rF]

# Build the coefficient matrix for equilibrium equations
coefficients = np.zeros((6, 6))

# Sum of forces in x, y, z directions
for idx, force_vec in enumerate(forces_list):
    coefficients[0, idx] = force_vec[0]  # Sum Fx
    coefficients[1, idx] = force_vec[1]  # Sum Fy
    coefficients[2, idx] = force_vec[2]  # Sum Fz

# Sum of moments about wheel center in x, y, z directions
for idx, (force_vec, r_vec) in enumerate(zip(forces_list, positions_list)):
    moment = np.cross(r_vec, force_vec)
    coefficients[3, idx] = moment[0]  # Sum Mx
    coefficients[4, idx] = moment[1]  # Sum My
    coefficients[5, idx] = moment[2]  # Sum Mz

# Right-hand side (negative of applied forces and moments)
rhs = -np.array([
    fap[0],  # Fx
    fap[1],  # Fy
    fap[2],  # Fz
    mom[0],  # Mx
    mom[1],  # My
    mom[2],  # Mz
])

# Solve the system of equations
forces = solve(coefficients, rhs)

# Print the forces in the suspension members
print("Forces in suspension members (N):")
member_names = ["F_A (Fore Bottom Arm)", "F_B (Aft Bottom Arm)",
                "F_C (Fore Top Arm)", "F_D (Aft Top Arm)", "F_E (Damper)", "F_F (Tie Rod)"]
for name, force_value in zip(member_names, forces):
    print(f"{name}: {force_value:.2f}")

for name, force_value in zip(member_names, forces):
    print(f"{force_value:.2f}")
# Compute the net force at the Lower Ball Joint (LBJ)
LBJ_force = forces[0] * forceA + forces[1] * forceB + forces[4] * forceE
print("\nLower Ball Joint Force (N):\n", LBJ_force)

# Compute the net force at the Upper Ball Joint (UBJ)
UBJ_force = forces[2] * forceC + forces[3] * forceD + forces[4] * forceE
print("\nUpper Ball Joint Force (N):\n", UBJ_force)

# Visualizing the A-Arm Geometry in 3D
def plot_link(ax, p1, p2, color, label):
    xs = [p1[0], p2[0]]
    ys = [p1[1], p2[1]]
    zs = [p1[2], p2[2]]
    ax.plot(xs, ys, zs, color=color, label=label)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Plot each suspension component
plot_link(ax, fb_upright, fb_chassis, 'green', 'Fore Bottom Arm')
plot_link(ax, ab_upright, ab_chassis, 'green', 'Aft Bottom Arm')
plot_link(ax, ft_upright, ft_chassis, 'blue', 'Fore Top Arm')
plot_link(ax, at_upright, at_chassis, 'blue', 'Aft Top Arm')
plot_link(ax, d_upright, d_chassis, 'red', 'Damper')
plot_link(ax, tr_upright, tr_chassis, 'purple', 'Tie Rod')

# Mark the wheel center and contact patch
ax.scatter(wheel_center[0], wheel_center[1], wheel_center[2], color='orange', label='Wheel Center')
ax.scatter(cP[0], cP[1], cP[2], color='black', label='Contact Patch')

# Set plot labels and aspect ratio
ax.set_xlabel('X (m)')
ax.set_ylabel('Y (m)')
ax.set_zlabel('Z (m)')
ax.set_title('Suspension Geometry')
ax.legend()
plt.show()

here's my script. 

Here's a test case: Let's say i manually overrode the lateral forces as:  fap = np.array([F_long_wheel, -2397, 1201])  # Applied force at contact patch (N)

This should result in something close in magnitude, to the following output array (obviously distributions iwll change due to geometry)

     793.416
 123.370
  -181.696 
    -46.814
  -148.082
    -253.756
 416.356

These forces are in lbf, please convert to N before you analyze the code. 

Using this, change any sign inconsistencies necessary, as well as correct any physics errors you see! Do not hard code the test case.