import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi

# Parameters
m=1     # mass kg
b=0.5   # drag coefficient
g=9.81  # acceleration due to gravity
te=60*pi/180  # angle
vo=100  # initial velocity m/s

# Initial conditions
x0=0
y0=0
vx0=vo*cos(te)  # initial velocity in x
vy0=vo*sin(te)  # initial velocity in y

# Time
t0=0
tf=1000
dt=0.001  # 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):
    # 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
    
    # 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:
            tm = t[i]
            xm = x[i]
            ymax = y[i]

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

    # Return the results
    return x[:i+1], y[:i+1], vx[:i+1], vy[:i+1], tm, tv, ymax, xm, xmax

# Solve the equations using the fourth-order Runge-Kutta method
x1, y1, vx, vy, tm, tv, ymax, xm, xmax=rk4(m,b,g,te,vo,dt,t)

# Plot the trajectories
plt.plot(x1,y1)
# Point of maximum height
plt.scatter(xm, ymax, color='red', s=25)  
plt.xlabel('$x(m)$', fontsize=12)
plt.ylabel('$y(m)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Print the results
print(f'Flight Time: {tv:.2f} seconds')
print(f'Maximum Range: {xmax:.2f} meters')
print(f'Time to Maximum Height: {tm:.2f} seconds')
print(f'Maximum Height: {ymax:.2f} meters')