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

# Parameters
belt_speed = 4  # meters per minute
oscillation_time = 3  # seconds per pass (forward and backward)
coating_height = 1.9  # meters
num_guns = 4  # 4 guns vertically
spacing_between_guns = coating_height / (num_guns - 1)  # Distance between guns
spray_width = 0.2  # meters (200 mm)
oscillation_distance = coating_height

# Calculate the speed of the oscillator
oscillator_speed = (2 * oscillation_distance) / oscillation_time  # meters per second
oscillator_speed_mpm = oscillator_speed * 60  # convert to meters per minute

# Time range for simulation
time_steps = np.linspace(0, 10, 200)
belt_positions = belt_speed / 60 * time_steps  # Belt positions in meters
oscillator_positions = (oscillator_speed * (time_steps % (oscillation_time / 2)))

# Adjust the oscillator position for the backward pass
oscillator_positions[time_steps % oscillation_time >= oscillation_time / 2] = (
    oscillation_distance - oscillator_positions[time_steps % oscillation_time >= oscillation_time / 2]
)

# Create 3D plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# Set plot limits
ax.set_xlim(0, 10)
ax.set_ylim(0, coating_height)
ax.set_zlim(0, 1)

# Labels
ax.set_xlabel('Time (seconds)')
ax.set_ylabel('Height of Coating Area (meters)')
ax.set_zlabel('Coverage Intensity')

# Initialize lines for the guns
lines = []
for _ in range(num_guns):
    line, = ax.plot([], [], [], lw=2)
    lines.append(line)

# Animation update function
def update(num, belt_positions, oscillator_positions, lines):
    coverage = np.zeros(int(coating_height / spray_width))
    covered_sections = int(oscillator_positions[num] / spray_width)
    for gun in range(num_guns):
        section_index = covered_sections + gun * spacing_between_guns / spray_width
        if section_index < coverage.shape[0]:
            coverage[int(section_index):int(section_index) + 1] += 1

    for i, line in enumerate(lines):
        gun_position = i * spacing_between_guns
        line.set_data(time_steps[:num], np.full(num, gun_position))
        line.set_3d_properties(coverage[:num])

    return lines

# Create animation
ani = animation.FuncAnimation(
    fig, update, frames=len(time_steps), fargs=(belt_positions, oscillator_positions, lines),
    interval=50, blit=False)

plt.show()