#!/usr/bin/env python3
"""
PSC6125 template:
Linear equatorial shallow-water model on an equatorial beta plane.

Equations:
    du/dt - y v = - dp/dx - r u
    dv/dt + y u = - dp/dy - r v
    dp/dt + du/dx + dv/dy = Q - eps p

Students can modify:
    1. Initial condition
    2. Forcing term Q(x,y,t)

Output:
    mp4 animation
    png final figure
"""

import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# ============================================================
# Parameters
# ============================================================
DT = 0.02
NT = 600

DX = 0.2
DY = 0.2

# Plot settings
QUIVER_X = 6
QUIVER_Y = 4
clevs = np.arange(-3.4, 3.8, 0.4)


# ============================================================
# Grid
# ============================================================
x = np.arange(-10 + DX/2, 10 - DX/2, DX)
y = np.arange(-5 - DY/2, 5 + DY/2, DY)

X, Y = np.meshgrid(x, y)
# Or Memory efficient way
#X = x[None, :]
#Y = y[:, None]


# ============================================================
# Math helper functions
# ============================================================
def hermiteH(n, y):
    """Physicists' Hermite polynomial H_n(y)."""
    if n == 0:
        return np.ones_like(y)
    if n == 1:
        return 2.0 * y

    Hnm1 = np.ones_like(y)
    Hn = 2.0 * y

    for k in range(1, n):
        Hnp1 = 2.0 * y * Hn - 2.0 * k * Hnm1
        Hnm1, Hn = Hn, Hnp1

    return Hn


def ddx(a):
    """Zonal derivative with cyclic boundary condition."""
    return (np.roll(a, -1, axis=1) - np.roll(a, 1, axis=1)) / (2.0 * DX)


def ddy(a):
    """Meridional derivative."""
    da = np.gradient(a, y, axis=0)
    da[0, :] = 0.0
    da[-1, :] = 0.0
    return da


# ============================================================
# Initial condition
# Students modify this section
# ============================================================
def initial_condition(X, Y):
    A = 2.5
    x0 = 0.0
    sigma_x = 3.0

    env_x = np.exp(-0.5 * ((X - x0) / sigma_x)**2)
    env_y = np.exp(-0.5 * Y**2)

    phase = 0.5 * (X - x0)

    # -- Kelvin Wave --
    #p = A * env_x * env_y * np.sin(phase)
    #u = p.copy()
    #v = np.zeros_like(p)

    # -- Zero initial conditions --
    A = 0.0
    p = A * env_x * env_y * np.sin(phase)
    u = np.zeros_like(p)
    v = np.zeros_like(p)

    return u, v, p


# ============================================================
# Forcing term
# Students modify this section
# ============================================================
def forcing_term(X, Y, t):
     Q0 = 2.5
     xQ = 0.0
     yQ = 0.0
     sigxQ = 0.75
     sigyQ = 0.75

     #--  Forcing  --
     #Q = Q0 * np.exp(-0.5 * ((X - xQ) / sigxQ)**2) \
     #       * np.exp(-0.5 * ((Y - yQ) / sigyQ)**2)

     #--  No forcing  --
     Q = np.zeros_like(X + Y)
     return Q


# ============================================================
# Model equation
# ============================================================
def model_step(u, v, p, t):
    """
    One forward-Euler time step.
    """

    # Momentum and Pressure Damping; use 0.08 for example 
    r = 0.00      # momentum damping
    eps = 0.00    # pressure damping

    Q = forcing_term(X, Y, t)

    dudt = Y * v - ddx(p) - r * u
    dvdt = -Y * u - ddy(p) - r * v
    dpdt = -(ddx(u) + ddy(v)) - eps * p + Q

    u_new = u + DT * dudt
    v_new = v + DT * dvdt
    p_new = p + DT * dpdt

    return u_new, v_new, p_new


# ============================================================
# Time integration
# ============================================================
u_save = np.zeros((NT, len(y), len(x)))
v_save = np.zeros_like(u_save)
p_save = np.zeros_like(u_save)
t_save = np.zeros(NT)

u, v, p = initial_condition(X, Y)
t = 0.0

u_save[0, :, :] = u
v_save[0, :, :] = v
p_save[0, :, :] = p
t_save[0] = t

for it in range(NT - 1):
    u, v, p = model_step(u, v, p, t)
    t = t + DT

    u_save[it + 1, :, :] = u
    v_save[it + 1, :, :] = v
    p_save[it + 1, :, :] = p
    t_save[it + 1] = t


# ============================================================
# Animation setup
# ============================================================
fig, ax = plt.subplots(figsize=(8, 4))

cf = ax.contourf(x, y, p_save[0], clevs, cmap="RdYlBu_r", extend="both")

xs = x[::QUIVER_X]
ys = y[::QUIVER_Y]
XXq, YYq = np.meshgrid(xs, ys)

q = ax.quiver(
    XXq, YYq,
    u_save[0, ::QUIVER_Y, ::QUIVER_X],
    v_save[0, ::QUIVER_Y, ::QUIVER_X],
    scale=30,
    pivot="mid",
)

cbar = fig.colorbar(cf, ax=ax, orientation="horizontal", pad=0.18)
cbar.set_label("Pressure / geopotential anomaly")

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title("Linear Equatorial Shallow-Water Model")


def update(frame):
    global cf

    cf.remove()

    cf = ax.contourf(
        x, y, p_save[frame],
        clevs,
        cmap="RdYlBu_r",
        extend="both",
        zorder=1,
    )

    q.set_UVC(
        u_save[frame, ::QUIVER_Y, ::QUIVER_X],
        v_save[frame, ::QUIVER_Y, ::QUIVER_X],
    )
    q.set_zorder(10)

    ax.set_title(f"Linear Equatorial Shallow-Water Model, t = {t_save[frame]:.2f}")

    return q,


ani = animation.FuncAnimation(
    fig,
    update,
    frames=NT,
    interval=30,
    blit=False,
)


# ============================================================
# Save animation and final figure
# ============================================================
plt.tight_layout()

script_name = os.path.splitext(os.path.basename(__file__))[0]

outfile_mp4 = script_name + ".mp4"
writer = animation.FFMpegWriter(fps=30, bitrate=1800)
ani.save(outfile_mp4, writer=writer)
print("Animation saved as:", outfile_mp4)

outfile_png = script_name + ".png"
plt.savefig(outfile_png, dpi=150)
print("Final figure saved as:", outfile_png)

plt.show()