#!/usr/bin/env python3
# Plot surface ocean Ekman spiral solution. See notes-friction.pdf
# EPS131, 2026, by ChatGPT from the corresponding Matlab scrip

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

plt.rcParams.update({"font.size": 14})

delta = 50  # meter
z = -(np.arange(0, 1500 + 0.1, 0.1))
A = 1

u = A * np.exp(z / delta) * (np.cos(z / delta) - np.sin(z / delta))
v = A * np.exp(z / delta) * (np.cos(z / delta) + np.sin(z / delta))

fig1 = plt.figure(1)

ax = fig1.add_subplot(1, 2, 1)
ax.plot(u, v)
ax.set_xlabel("u")
ax.set_ylabel("v")
ax.set_title("Ekman spiral, zoom in to see details")
ax.text(0.75, 0.9, "z=0", transform=ax.transAxes)
ax.text(0.025, 0.2, r"z=-$\infty$", transform=ax.transAxes)

ax.set_xlim([-0.25, 1.1])
ax.set_ylim([-0.25, 1.1])
ax.set_aspect("equal", adjustable="box")

# --- subplot(1,2,2)
ax = fig1.add_subplot(1, 2, 2)
ax.plot(u, z, "r")
ax.plot(v, z, "b")
ax.set_ylim([-300, 0])

ax.set_xlabel("u,v")
ax.set_ylabel("z")
ax.legend(["u", "v"], loc="lower right")
ax.set_title("u,v")

plt.tight_layout()
fig1.set_size_inches(8, 4)

if True:
    fig1.savefig(os.path.join("Output", "my_Ekman_spiral.pdf"), bbox_inches="tight")

plt.show()