# Eksempel 6.2.1 AMS side 153
# Mars Express <3
# Her beregner vi hastighet og akselerasjon
# Ann-Cecilie Larsen, 11 March 2023
# a.c.larsen@fys.uio.no
# Last version: 04 March 2024

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


# Lese datafil (h?y oppl?sning)
t, x, y, z = np.loadtxt('marsexpresshr.txt', usecols = [0, 1, 2, 3], unpack = True)
n = len(t)

r = np.zeros((n,3),float) # Posisjonsvektor, antall rader: n; antall kolonner: 3
r[:, 0] = x # For alle rader fyller vi x-verdiene i kolonne 0
r[:, 1] = y # For alle rader fyller vi y-verdiene i kolonne 1
r[:, 2] = z # For alle rader fyller vi z-verdiene i kolonne 2

v = np.zeros((n,3),float) # Hastighetsvektor, antall rader: n; antall kolonner: 3
a = np.zeros((n,3),float) # Akselerasjonsvektor, antall rader: n; antall kolonner: 3

v_abs = np.zeros((n,1),float) # Lengden til hastighetsvektor, v = \sqrt{v_x^2 + v_y^2 + v_z^2}
a_abs = np.zeros((n,1),float) # Lengden til akselerasjonsvektor, a = \sqrt{a_x^2 + a_y^2 + a_z^2}

# Vi regner ut hastigheten, den deriverte av posisjonen. 
# Vi f?r n-1 punkter for hastigheten
for i in range(0, n-1):
	dt = t[i+1]-t[i] # regner ut tidsintervallet
	v[i, :] = (r[i+1, :] - r[i, :])/dt # Hastigheten
	v_abs[i] = np.linalg.norm(v[i, :]) # Farten

# S? regner vi ut endring i hastighet per tidsintervall, 
# deriverer hastighet og f?r akselerasjonen. 
# MERK at l?kka starter i 1 og stopper i n-1! 
# Dette fordi vi f?r n-2 punkter for akselerasjonen.
for i in range(1, n-1):
	dt = t[i]-t[i-1] # regner ut tidsintervallet
	a[i, :] = (v[i, :] - v[i-1, :])/dt # Akselerasjonen
	a_abs[i] = np.linalg.norm(a[i, :]) # St?rrelsen p? akselerasjonen

# Plotte x- og y-posisjon
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
plt.plot(r[:, 0], r[:, 1]) 
plt.xlabel('x [km]')
plt.ylabel('y [km]')
plt.axis('equal')
fig.tight_layout()
fig.savefig('highres_position.png')

# Plotte farten
fig2 = plt.figure() 
ax2  = fig2.add_subplot(2, 1, 1)
plt.plot(t[0: n-1],v_abs[0: n-1]) 
plt.xlabel('t [dager]')
plt.ylabel('v [km/dag]')
#plt.ylim([2.0E+06, 4.0E+06])

# Plotte str. p? akselerasjonen
ax2  = fig2.add_subplot(2, 1, 2)
plt.plot(t[1: n-5],a_abs[1: n-5]) 
plt.xlabel('t [dager]')
plt.ylabel('a [km/dag$^2$]')
fig2.tight_layout()
fig2.savefig('highres_v_and_a.png')

# Plotte posisjonen i 3D
fig3 = plt.figure() 
ax3 = plt.axes(projection = '3d')
ax3.plot3D(r[:, 0], r[:, 1], r[:, 2], 'gray')
ax3.set_xlabel('x [km]')
ax3.set_ylabel('y [km]')
ax3.set_zlabel('z [km]')
plt.axis('equal')
fig3.tight_layout()
fig3.savefig('position3D.png')

# Plotte alle akselerasjonskomponentene
fig4 = plt.figure() 
ax4  = fig4.add_subplot(3, 1, 1)
plt.plot(t[1: n-2],a[1: n-2, 0]) # Plotte x-komp av akselerasjonen
plt.xlabel('t [dager]')
plt.ylabel('a$_x$ [km/dag$^2$]')

ax4  = fig4.add_subplot(3, 1, 2)
plt.plot(t[1: n-2],a[1: n-2, 1]) # Plotte y-komp av akselerasjonen
plt.xlabel('t [dager]')
plt.ylabel('a$_y$ [km/dag$^2$]')

ax4  = fig4.add_subplot(3, 1, 3)
plt.plot(t[1: n-2],a[1: n-2, 2]) # Plotte z-komp av akselerasjonen
plt.xlabel('t [dager]')
plt.ylabel('a$_z$ [km/dag$^2$]')
fig4.tight_layout()
fig4.savefig('akselerasjon.png')


plt.show()