#To run the program with greater precision, for example with dt=0.0001, 
#it is recommended to run this file in JupyterLab.
import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi

# Parameters
m = 1     # mass kg
b = 0.1   # drag coefficient
g = 9.81  # acceleration due to gravity
vo = 700  # initial velocity m/s

# Time
t0 = 0
tf = 1000
dt = 0.01  # time step
t = np.arange(t0, tf, dt)

# Function to calculate accelerations
def ac(vx, vy, m, b, g):
    ax = -b * vx * np.sqrt(vx**2 + vy**2) / m
    ay = -g - b * vy * np.sqrt(vx**2 + vy**2) / m
    return ax, ay

def rk4(m, b, g, te, vo, dt, t):
    # Initial conditions
    x0 = 0
    y0 = 0
    vx0 = vo * cos(te)  # initial velocity in x
    vy0 = vo * sin(te)  # initial velocity in y
    
    # Arrays to store solutions
    x = np.zeros_like(t)
    y = np.zeros_like(t)
    vx = np.zeros_like(t)
    vy = np.zeros_like(t)

    # Seeds
    x[0] = x0
    y[0] = y0
    vx[0] = vx0  # initial velocity in x
    vy[0] = vy0  # initial velocity in y
    
    ymax = y0
    xm = x0
    # Solve the equations using the fourth-order Runge-Kutta method
    for i in range(1, len(t)):
        # Calculate intermediate increments for vx and vy
        ax1, ay1 = ac(vx[i-1], vy[i-1], m, b, g)
        k1vx = ax1 * dt
        k1vy = ay1 * dt
        k1x = vx[i-1] * dt
        k1y = vy[i-1] * dt

        ax2, ay2 = ac(vx[i-1] + 0.5 * k1vx, vy[i-1] + 0.5 * k1vy, m, b, g)
        k2vx = ax2 * dt
        k2vy = ay2 * dt
        k2x = (vx[i-1] + 0.5 * k1vx) * dt
        k2y = (vy[i-1] + 0.5 * k1vy) * dt

        ax3, ay3 = ac(vx[i-1] + 0.5 * k2vx, vy[i-1] + 0.5 * k2vy, m, b, g)
        k3vx = ax3 * dt
        k3vy = ay3 * dt
        k3x = (vx[i-1] + 0.5 * k2vx) * dt
        k3y = (vy[i-1] + 0.5 * k2vy) * dt

        ax4, ay4 = ac(vx[i-1] + k3vx, vy[i-1] + k3vy, m, b, g)
        k4vx = ax4 * dt
        k4vy = ay4 * dt
        k4x = (vx[i-1] + k3vx) * dt
        k4y = (vy[i-1] + k3vy) * dt

        # Update velocities and positions
        vx[i] = vx[i-1] + (k1vx + 2*k2vx + 2*k3vx + k4vx) / 6
        vy[i] = vy[i-1] + (k1vy + 2*k2vy + 2*k3vy + k4vy) / 6
        x[i] = x[i-1] + (k1x + 2*k2x + 2*k3x + k4x) / 6
        y[i] = y[i-1] + (k1y + 2*k2y + 2*k3y + k4y) / 6

        # Check if it reached maximum height
        if vy[i-1] > 0 and vy[i] <= 0:
            ymax = y[i]
            xm = x[i]

        # Check if it reached the ground
        if y[i] <= 0:
            break

    # Return the results
    return x[:i+1], y[:i+1], ymax, xm

# Plot the trajectories for angles from 0° to 90°
plt.figure(figsize=(6,3))
for angle in range(0, 181, 5):
    te = angle * pi / 180  # Convert angle to radians
    x1, y1, ymax, xm = rk4(m, b, g, te, vo, dt, t)
    plt.plot(x1, y1, color='blue', label=f'{angle}°',linewidth=0.7)  # Light green color
    plt.scatter(xm, ymax, color='red', s=8)  # Mark the point of maximum height

plt.xlabel('$x(m)$', fontsize=12)
plt.ylabel('$y(m)$', fontsize=12)
plt.grid(True)
plt.gca().set_aspect('equal', adjustable='box')  # Set aspect ratio to 1:1
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.axis('equal')
plt.tight_layout()
plt.show()