#%% 0. Imports
import numpy as np
import matplotlib.pyplot as plt

#%% 1. Condições iniciais
lat = 16.0  # Graus Norte
lon = -24.0 # Graus Este

wind = 1             # Velocidade módulo exemplo para os ventos alíseos superficiais (m/s)
u = -wind            # Velocidade zonal inicial
v = 0                # Velocidade meridional inicial

dt = 3600               # Timestep de uma hora (em segundos)
days = 1.5 
t_end = days * 24 * 3600   # Tempo total de integração

# Parâmetros do planeta
Omega = 7.292e-5        # Velocidade de rotação
R = 6.371e6             # Raio do planeta

# Listas para guardar as trajetórias
lats = [lat]
lons = [lon]

#%% 2. Integração numérica
for t in np.arange(0, t_end, dt):
    
    # Converter a latitude para radianos
    phi = np.deg2rad(lats[-1])
    
    # Calcular o parâmetro de Coriolis
    f = 2 * Omega * np.sin(phi)

    # Aplicar a aceleração de Coriolis: 
    # du/dt =  f * v_n
    # dv_n/dt = -f * u
    du =  f * v * dt
    dv = -f * u  * dt

    u += du
    v += dv
    
    # Keep speed constant but adjust direction
    speed = np.sqrt(u**2 + v**2)
    u = wind * (u / speed)
    v = wind * (v / speed)

    # Convert velocity components into changes in lat/lon
    dlat = (v * dt) / R
    dlon = (u * dt) / (R * np.cos(phi))
    
    lats.append(np.degrees(dlat)+lats[-1])
    lons.append(np.degrees(dlon)+lons[-1])

#%% 3. Mostrar trajetória
plt.figure(figsize=(8,8))
plt.grid()

plt.scatter(lon,lat,c='r')
plt.plot(lons, lats)
plt.title("Trajetória da partícula")
plt.xlabel("Longitude")
plt.ylabel("Latitude")