import numpy as np
import matplotlib.pyplot as plt

C = 0.47
m = 4.08
rho_a = 1.21
r_b = 0.1016 
g = 9.81
alpha = 0.5 * rho_a * np.pi * r_b * r_b * C / m


# Question 1: the rewriting is necessary because the RK4 only accepts first-order ODEs.


# Question 2: In the absence of friction, the cannonball hits the ground after a time t=2*v*sin(theta)/g, with friction, the fall is faster (because energy decays).
def F(Z):
    return np.array([Z[2], Z[3], - alpha * Z[2] * np.sqrt(Z[2] ** 2 + Z[3] ** 2), - g - alpha * Z[3] * np.sqrt(Z[2] ** 2 + Z[3] ** 2)])


def rk4(h, vi, thetai, xi=0., yi=0., ti=0.):
    thetai *= np.pi / 180.
    Z = np.array([xi, yi, vi * np.cos(thetai), vi * np.sin(thetai)])
    tf = 2. * Z[-1] / g
    t_arr = np.arange(ti, tf + h, h)
    coordinates = np.zeros((t_arr.size, 2))
    coordinates[0] = Z[:2]
    for i in range(1, t_arr.size):
        K1 = F(Z) * h
        K2 = F(Z + K1 / 2.) * h
        K3 = F(Z + K2 / 2.) * h
        K4 = F(Z + K3) * h
        Z += (K1 + 2. * K2 + 2. * K3 + K4) / 6.
        coordinates[i] = Z[:2]
        if Z[1] < 0:
            break
    x_arr = coordinates[:i,0]
    y_arr = coordinates[:i,1]
    return t_arr, x_arr, y_arr
    

# Question 3:
h = 0.01
vi = 700.
theta_val = np.arange(10., 90., 10.)
plt.figure()
for theta in theta_val:
    t, x, y = rk4(h, vi, theta)
    plt.plot(x, y, label=r'$\theta_\mathrm{i}=$' + str(theta))
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.title(r'Trajectory of the cannonball for $v_\mathrm{i}=$' + str(vi) + r' m/s and different initial angles $\theta_\mathrm{i}$')
plt.grid()
plt.legend()
plt.show()


# Question 4a/4b: h = 0.001 is necessary to get the tolerance of 0.1m. There are two possible couples of (theta_min, theta_max) which are (0., 10.) and (50., 60.)
h = 0.001
vi = 700.
xt = 1000.
yt = 15.


def extract_y_impact(x, y):  # Give the value of y at xt if the cannonball has reached xt, and 0 otherwise.
    if x[-1] > xt:
       return y[np.argmin(np.absolute(x - xt))]
    else:
       return 0.

plt.figure()
plt.plot(xt, yt, 'ok')
theta_i = []
for thetas in [(1e-5, 10.), (50., 60.)]:
    theta_min, theta_max = thetas
    t, x, y = rk4(h, vi, theta_min)
    ymin = extract_y_impact(x, y)
    t, x, y = rk4(h, vi, theta_max)
    ymax = extract_y_impact(x, y)
    theta_mid = 0.5 * (theta_min + theta_max)
    t, x, y = rk4(h, vi, theta_mid)
    ymid = extract_y_impact(x, y)
    while np.abs(theta_max - theta_min) > 1e-5:
        if (ymid - yt) * (ymin - yt) > 0.:
            theta_min = theta_mid
            ymin = ymid
        else:
            theta_max = theta_mid
            ymax = ymid
        theta_mid = 0.5 * (theta_min + theta_max)
        t, x, y = rk4(h, vi, theta_mid)
        ymid = extract_y_impact(x, y)
    theta_i += [theta_mid]
    plt.plot(x, y, label=r'$\theta_\mathrm{i}=$' + str(theta_mid) + r'$^\circ$')
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.grid()
plt.title(r'Optimal trajectories to reach the target')
plt.legend()
plt.show()
       
       
# Question 4c: The first trajectory has the highest kinetic energy at impact, better for war!
def rk4_with_vel(h, vi, thetai, xi=0., yi=0., ti=0.):
    thetai *= np.pi / 180.
    Z = np.array([xi, yi, vi * np.cos(thetai), vi * np.sin(thetai)])
    tf = 2. * Z[-1] / g
    t_arr = np.arange(ti, tf + h, h)
    coordinates = np.zeros((t_arr.size, 4))
    coordinates[0] = Z
    for i in range(1, t_arr.size):
        K1 = F(Z) * h
        K2 = F(Z + K1 / 2.) * h
        K3 = F(Z + K2 / 2.) * h
        K4 = F(Z + K3) * h
        Z += (K1 + 2. * K2 + 2. * K3 + K4) / 6.
        coordinates[i] = Z
        if Z[1] < 0:
            break
    x_arr = coordinates[:i,0]
    y_arr = coordinates[:i,1]
    vx_arr = coordinates[:i,2]
    vy_arr = coordinates[:i,3]
    return t_arr, x_arr, y_arr, vx_arr, vy_arr
    
for theta in theta_i:
    t, x, y, vx, vy = rk4_with_vel(h, vi, theta)
    j = np.argmin(np.absolute(y - yt))
    print('Kinetic energy = ', 0.5 * m * (vx[j] ** 2 + vy[j] ** 2), ' J')