import time
import random 
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D

# --- CONFIGURATION ---
ball_radius = 3
box_size = 200
particles_per_side = 7  # 7x7x7 = 343 particles
total_particles = particles_per_side**3
frame_limit = 500

class Particle:
    def __init__(self, x, y, z, dx, dy, dz):
        self.x, self.y, self.z = x, y, z
        self.dx, self.dy, self.dz = dx, dy, dz

    def get_pos(self):
        return np.array([self.x, self.y, self.z])

    def get_vel(self):
        return np.array([self.dx, self.dy, self.dz])

    def get_distance(self, other):
        return np.sqrt((self.x - other.x)**2 + (self.y - other.y)**2 + (self.z - other.z)**2)

    def handle_collision(self, other):
        # 3D Elastic Collision Physics (Vector Form)
        v1, v2 = self.get_vel(), other.get_vel()
        x1, x2 = self.get_pos(), other.get_pos()
        
        # Center of mass frame calculation
        v1_new = v1 - (v1 - v2).dot(x1 - x2) / (np.linalg.norm(x1 - x2)**2) * (x1 - x2)
        v2_new = v2 - (v2 - v1).dot(x2 - x1) / (np.linalg.norm(x2 - x1)**2) * (x2 - x1)
        
        self.dx, self.dy, self.dz = v1_new
        other.dx, other.dy, other.dz = v2_new

# --- INITIALIZATION ---
particles = []
# Set system temperature (Total Energy)
target_total_energy = total_particles * 1
initial_energies = np.random.uniform(0.1, 1.5, total_particles)
initial_energies = initial_energies * target_total_energy / initial_energies.sum()
velocity_magnitudes = np.sqrt(2 * initial_energies)

# Arrange particles in a 3D grid with random directions
k = 0
grid_spacing = box_size / (particles_per_side + 1)
for i in range(1, particles_per_side + 1):
    for j in range(1, particles_per_side + 1):
        for m in range(1, particles_per_side + 1):
            # Random direction in 3D using spherical coordinates
            phi = random.uniform(0, 2 * np.pi)
            cos_theta = random.uniform(-1, 1)
            theta = np.arccos(cos_theta)
            
            dx = velocity_magnitudes[k] * np.sin(theta) * np.cos(phi)
            dy = velocity_magnitudes[k] * np.sin(theta) * np.sin(phi)
            dz = velocity_magnitudes[k] * np.cos(theta)
            
            particles.append(Particle(i * grid_spacing, j * grid_spacing, m * grid_spacing, dx, dy, dz))
            k += 1

# Data history for playback
pos_history = []
vel_history = []

print("Simulating 3D environment...")
for f in range(frame_limit):
    frame_velocities = []
    frame_positions = [[], [], []]
    
    # 1. Inter-particle collisions
    for i in range(len(particles)):
        for j in range(i + 1, len(particles)):
            if particles[i].get_distance(particles[j]) < ball_radius * 2:
                particles[i].handle_collision(particles[j])
    
    # 2. Movement and Wall collisions
    for p in particles:
        p.x += p.dx; p.y += p.dy; p.z += p.dz
        
        # Bounce logic for each axis
        for axis_name in ['x', 'y', 'z']:
            pos_val = getattr(p, axis_name)
            vel_val = getattr(p, 'd' + axis_name)
            if pos_val > box_size - ball_radius or pos_val < ball_radius:
                setattr(p, 'd' + axis_name, vel_val * -1)
        
        frame_positions[0].append(p.x)
        frame_positions[1].append(p.y)
        frame_positions[2].append(p.z)
        frame_velocities.append(np.sqrt(p.dx**2 + p.dy**2 + p.dz**2))
        
    pos_history.append(frame_positions)
    vel_history.append(frame_velocities)
    if f % 20 == 0: print(f"Processing frame {f}...")

# --- THEORETICAL MAXWELL-BOLTZMANN CURVE (3D) ---
v_axis = np.linspace(0, 4, 100)
# Formula: f(v) = (m/2piKT)^1.5 * 4pi * v^2 * exp(-mv^2 / 2KT)
# We simplify based on simulated energy units
a = 0.75 
theoretical_mb = (4 / np.sqrt(np.pi)) * (a**1.5) * (v_axis**2) * np.exp(-a * v_axis**2)

# --- ANIMATION ---
fig = plt.figure(figsize=(12, 6))
ax_3d = fig.add_subplot(121, projection='3d')
ax_hist = fig.add_subplot(122)

def animate(frame_idx):
    ax_3d.clear()
    ax_hist.clear()
    
    # Left: 3D Visualization
    pos = pos_history[frame_idx]
    ax_3d.scatter(pos[0], pos[1], pos[2], s=12, c='crimson', edgecolors='black', alpha=0.6)
    ax_3d.set_xlim(0, box_size); ax_3d.set_ylim(0, box_size); ax_3d.set_zlim(0, box_size)
    ax_3d.set_title("3D Molecular Dynamics Simulation")
    
    # Right: Maxwell-Boltzmann Distribution
    current_vels = vel_history[frame_idx]
    ax_hist.hist(current_vels, bins=20, density=True, color='lightsteelblue', edgecolor='white', label='Simulated Data')
    ax_hist.plot(v_axis, theoretical_mb, 'r-', lw=2, label='Theoretical 3D M-B')
    
    ax_hist.set_xlim(0, 4); ax_hist.set_ylim(0, 1.2)
    ax_hist.set_title(f"Velocity Distribution (Iteration: {frame_idx})")
    ax_hist.set_xlabel("Speed (v)")
    ax_hist.set_ylabel("Probability Density")
    ax_hist.legend()

simulation = animation.FuncAnimation(fig, animate, frames=frame_limit, interval=50, repeat=True)
plt.tight_layout()
plt.show()