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)
    r=np.zeros_like(t)
    K=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
    K[0]=m*(vx0**2+vy0**2)/2
    
    # 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

        K[i]=m*(vx[i]**2+vy[i]**2)/2
        r[i]=np.sqrt(x[i]**2+y[i]**2)

        # 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

    # Find the point of minimum kinetic energy
    Km = np.argmin(K[:i+1])
    Kmin = K[Km]
    tKm = t[Km]
    
    # Height and horizontal range at the moment of minimum kinetic energy
    id = np.argmin(np.abs(t[:i+1]-tKm)) # Find the time closest to tKm
    yminK = y[id]
    xminK = x[id]

    # Return the additional results
    return x[:i+1], y[:i+1], vx[:i+1], vy[:i+1], K[:i], r[:i], t[:i], tm, tv, ymax, xm, xmax, Kmin, tKm, yminK, xminK

# Solve the equations using the fourth-order Runge-Kutta method
x, y, vx, vy, K, r, t, tm, tv, ymax, xm, xmax, Kmin, tKm, yminK, xminK=rk4(m,b,g,te,vo,dt,t)

# Plot of the trajectory
plt.plot(x,y)
# Point of maximum height
plt.scatter(xm, ymax, color='red', s=25)
plt.scatter(xminK, yminK, color='black', 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()

# First subplot: Plot of kinetic energy as a function of time
plt.figure(figsize=(14, 6))
plt.subplot(1, 2, 1)
plt.plot(t, K)
plt.scatter(tKm, Kmin, color='red', s=25)
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$K(J)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

# Second subplot: Plot of kinetic energy near the point of minimum energy
plt.subplot(1, 2, 2)
start = np.searchsorted(t, 0.2)
end = np.searchsorted(t, tv)
plt.plot(t[start:end], K[start:end])
plt.scatter(tKm, Kmin, color='red', s=25)
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$K(J)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Plot of the radial vector as a function of time
plt.figure(figsize=(7,6))
plt.plot(t,r)  
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$r(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')
print(f'Time to Minimum Kinetic Energy: {tKm:.2f} seconds')
print(f'Minimum Kinetic Energy: {Kmin:.2f} Joules')
print(f'Height at Minimum Kinetic Energy: {yminK:.2f} meters')
print(f'Horizontal Range at Minimum Kinetic Energy: {xminK:.2f} meters')