# FYS1100, vår 2024 og vår 2025
# Eksempel med en bordtennisball som spretter på gulvet,  
# Vi ser *ikke* bort fra luftmotstand her, noe som fører til
# at vi får to koblede diff-ligninger for akselerasjonen i x- og y-retning.
# I tillegg modellerer vi normalkraften fra gulvet på ballen som en fjærkraft.
# Vi har følgende krefter som virker (alle er vektorer):
# Luftmotstand F_D = -D |v| v
# Normalkraft N = k(R-y(t)), y(t) < R hvor R = ballens radius,
# eller       N = 0, y(t) >= R
# Gravitasjon G = -mg
# Vi ser her bort fra friksjon mellom gulv og ball.
# Konstanter er tatt fra Andreas Goergens pingpong.py fra FYS-MEK1110, vår 2020
# Ann-Cecilie Larsen, 11 Feb 2020
# Last update: 11 Mar 2024 (takk til Erlend Aune for oppryddingstips!)
# Nå er kreftene også numpy arrays, og y-komponenten av kreftene plottes også.
# a.c.larsen@fys.uio.no

# Importer nyttige Python-bibliotek
import numpy as np
import matplotlib.pyplot as plt

# Definerer konstanter og variabler vi vil trenge
# Jeg har stort sett brukt  konstantene fra pingpong.py 
g = 9.81 	# tyngdeakselerasjon i m/s^2
m = 0.0027	# masse til ballen i kg
v_T = 5.0	# terminalhastigheten i m/s.
D = m*g/(v_T**2)	# luftmotstandskoeffisienten i kg/m
R = 0.02	# radius til ballen i m (2 cm)
k = 100.0	# fjærkonstanten i N/m
x0 = 0	# initiell posisjon i x-retning, definert til å være i 0 m
y0 = 0.8	# initiell høyde (y-retning). Vi kan teste forskjellige start-høyder om vi vil
r0 = np.array([x0, y0])	# initiell r-vektor, med sin x- og y-komponent
# Her lager vi start-hastigheten ved t0 = 0 s
v0 = np.array([1.5, -0.5])  
time = 5	# maximum tid vi ser på i sekunder 
dt = 0.0001	# tidssteg i sekunder - pass på at det er lite nok!

# Her bruker vi numpys ceil-funksjon for å bestemme antall elementer vi vil ha i vektorene og matrisene.
# "The ceil of a scalar b is the smallest integer i such that i>= b"
# https://www.geeksforgeeks.org/numpy-ceil-python/
n = int(np.ceil(time/dt))

# Nå definerer vi en vektor for tiden t,  og matriser for posisjonen r (x- og y-komponent ved tid t), 
# hastigheten v (x- og y-komponent ved tid t), og akselerasjonen a (x- og y-komponent ved tid t)
# n rader, to kolonner:
# r = [x0, y0]
#     [x1, y1]
#		...
#	  [xn, yn]

t = np.zeros(n)
r = np.zeros((n, 2))
# For å sjekke hvordan r ser ut:
#print(r)
v = np.zeros((n, 2))
a = np.zeros((n, 2))

Fnet = np.zeros((n, 2))	# netto kraft på ballen
F_D = np.zeros((n, 2))	# luftmotstanden på ballen. Antar turbulente forhold
F_N = np.zeros((n, 2))	# normalkraften på ballen. Modelert som en fjærkraft
F_G = np.zeros((n, 2))	# gravitasjonskraften på ballen


# Nå bruker vi initialbetingelsene. 
# Denne litt kryptiske notasjonen betyr:
# "0": rad 0 -> første rad
# ":": "slice notation", her betyr det bare at vi tar alle kolonner i rad 0
# som er x0, y0 for posisjonen, og vx0, vy0 for hastigheten.
r[0, :] = r0 # Startposisjonen til r-vektoren
t[0] = 0	# Denne er forsåvidt satt allerede - vi har jo fylt hele vektoren med nuller da vi definerte den!
v[0, :] = v0	# Initialhastighet
a[0, :] = 0	# initial-akselerasjon

# Nå er vi klare for å kjøre ei løkke og regne ut alt vi vil ha:
for i in range (0,n-1):
	# Først regner vi ut normalkraften. 
	# r[i, 1] er y-komponenten til posisjonen ved rad ("tid") i
	if r[i, 1]<R:
		F_N[i, :] = k*(R-r[i, 1])*np.array([0, 1]) # Enhetsvektoren i y-retning = np.array([0,1])
		#print(r[i, 1])
	else: 
		F_N[i, :] = np.array([0, 0])
	F_D[i, :] = -D*np.linalg.norm(v[i, :])*v[i, :]
	F_G[i, :] = -m*g*np.array([0, 1]) # Enhetsvektoren i y-retning = np.array([0,1])
	Fnet[i, :] = F_N[i, :] + F_D[i, :] + F_G[i, :] 	# Nettokraften. Fortegnene er satt i linjene over, der F_D pg F_G er regnet ut.
	#Fnet[i, :] = F_N[i, :] + F_G[i, :] 	# Nettokraften uten luftmotstand
	a[i, :] = Fnet[i, :]/m 	    # akselerasjonen
	v[i+1, :] = v[i, :] + a[i, :]*dt # hastigheten
	# Vi bruker Euler-Cromer-metoden for å beregne posisjonsvektoren
	r[i+1, :] = r[i, :] + v[i+1, :]*dt 
	t[i+1] = t[i] + dt # her inkrementerer vi tiden


# Funksjonen subplots er veldig kjekk!
# Se https://matplotlib.org/api/_as_gen/matplotlib.pyplot.subplots.html
fig, ax = plt.subplots(nrows=3, figsize=(6, 7))
# r[:i, 0]: Første kolonne (0) som da er x-komponenten, for alle tider 0->i
# r[:i, 1]: Andre kolonne (1) som da er y-komponenten, for alle tider 0->i
# Vi har laget en matrise med n rader, en for hver tid, 
# og 2 kolonner, en for x-komponenten og en for y-komponenten
ax[0].plot(r[:i, 0], r[:i, 1])
ax[0].set_xlabel("Posisjon x [m]")
ax[0].set_ylabel("Posisjon y [m]")

# Her plotter vi de to hastighetskomponentene for alle tider
ax[1].plot(t[:i], v[:i, 0])
ax[1].plot(t[:i], v[:i, 1])
ax[1].set_xlabel('Tid t [s]')
ax[1].set_ylabel('v$_x$, v$_y$ [m/s]')
#ax[1].set_title('Hastighetskomponenter v$_x$ og v$_y$')
# Putte på en boks med info om hva grafene er for noe
ax[1].legend(["v$_x$", "v$_y$"])

# Her plotter vi de to akselerasjonskomponentene for alle tider
ax[2].plot(t[:i], a[:i, 0])
ax[2].plot(t[:i], a[:i, 1]) #,"." for markers
ax[2].set_xlabel('Tid t [s]')
ax[2].set_ylabel('a$_x$, a$_y$ [m/s$^2$]')
#ax[2].set_title('Akselerasjonskomponenter a$_x$ og a$_y$')
# Putte på en boks med info om hva grafene er for noe
ax[2].legend(["a$_x$", "a$_y$"])
fig.tight_layout()

# Lage en figur som viser kreftene som funksjon av tid
# y-komponenten er mest interessant, så plotter den her.
fig2, ax2 = plt.subplots(nrows=1, figsize=(6, 4))
ax2.plot(t[:i], F_N[:i, 1])
ax2.plot(t[:i], F_D[:i, 1])
ax2.plot(t[:i], F_G[:i, 1])
ax2.set_xlabel('Tid t [s]')
ax2.set_ylabel('y-komponent av kreftene [N]')
#ax[1].set_title('Hastighetskomponenter v$_x$ og v$_y$')
# Putte på en boks med info om hva grafene er for noe
ax2.legend(["F$_N$", "F$_D$", "F$_G$"])
fig2.tight_layout()
fig2.savefig('krefter_pingpongball.png')

plt.show()