import sympy as sp
import numpy as np

# Definicja danych wejściowych
r1 = 1.0   # średnica piłki 1
m1 = 1.0   # masa piłki 1
x1, y1, z1 = 0.0, 0.0, 0.0  # początkowe położenie piłki 1
vx1, vy1, vz1 = 1.0, 1.0, 1.0  # prędkość piłki 1

r2 = 1.0   # średnica piłki 2
m2 = 1.0   # masa piłki 2
x2, y2, z2 = 1.0, 1.0, 1.0  # początkowe położenie piłki 2
vx2, vy2, vz2 = -1.0, -1.0, -1.0  # prędkość piłki 2

# Promienie piłek
R1 = r1 / 2
R2 = r2 / 2

# Promień sumaryczny do sprawdzenia kolizji
R = R1 + R2

# Zmienne czasowe
t = sp.symbols('t')

# Położenie piłek w funkcji czasu
x1t = x1 + vx1 * t
y1t = y1 + vy1 * t
z1t = z1 + vz1 * t

x2t = x2 + vx2 * t
y2t = y2 + vy2 * t
z2t = z2 + vz2 * t

# Odległość między piłkami w funkcji czasu
distance = sp.sqrt((x2t - x1t)**2 + (y2t - y1t)**2 + (z2t - z1t)**2)

# Równanie kolizji
collision_eq = sp.Eq(distance, R)

# Rozwiązanie równania kolizji
collision_time = sp.solve(collision_eq, t)
collision_time = [t for t in collision_time if t >= 0]

# Sprawdzenie, czy nastąpi kolizja
if collision_time:
    t_collision = collision_time[0]
    
    # Punkt zderzenia
    x_collision = x1 + vx1 * t_collision
    y_collision = y1 + vy1 * t_collision
    z_collision = z1 + vz1 * t_collision
    
    collision_point = (x_collision.evalf(), y_collision.evalf(), z_collision.evalf())
else:
    collision_point = None

if collision_point:
    # Przeliczenie położenia w momencie kolizji
    x1_collision = x1 + vx1 * t_collision
    y1_collision = y1 + vy1 * t_collision
    z1_collision = z1 + vz1 * t_collision

    x2_collision = x2 + vx2 * t_collision
    y2_collision = y2 + vy2 * t_collision
    z2_collision = z2 + vz2 * t_collision

    # Wektor normalny w momencie zderzenia
    n = np.array([x2_collision - x1_collision, y2_collision - y1_collision, z2_collision - z1_collision]).astype(float)
    n = n / np.linalg.norm(n)

    # Wektory prędkości
    v1 = np.array([vx1, vy1, vz1])
    v2 = np.array([vx2, vy2, vz2])

    # Prędkości po zderzeniu
    v1_prime = v1 - (2 * m2 / (m1 + m2)) * np.dot(v1 - v2, n) * n
    v2_prime = v2 - (2 * m1 / (m1 + m2)) * np.dot(v2 - v1, n) * n

    new_velocities = (v1_prime, v2_prime)
else:
    new_velocities = None

print("Punkt kolizji:", collision_point)
print("="*20)
print("Nowa prędkość po zderzeniu:", new_velocities) 

--------------------------------------------------------------
Punkt kolizji: (0.211324865405187, 0.211324865405187, 0.211324865405187)
====================
Nowa prędkość po zderzeniu: (array([-1., -1., -1.]), array([1., 1., 1.]))