"""Waimea Bay, driven by the live NDBC 51201 buoy, through the surf kernel.

    python waimea.py                # live sea state off the bay mouth
    python waimea.py --winter       # a representative Eddie-day swell
    python waimea.py --hs 4 --tp 14 --frames 90

Buoy 51201 (Waimea Bay, HI) sits in ~200 m of water outside the bay and reports
swell and wind-wave trains separately, so the wavemaker below drives both: a
long-period groundswell plus the shorter local windswell riding on it. The North
Shore is seasonal. From roughly April to October 51201 reads 1-2 m and Waimea is
a swimming beach; the bay only does what it is famous for when the winter North
Pacific sends 6 m at 17 s.

The bathymetry here is an idealized Waimea: the real bay's dimensions (a ~600 m
mouth, the deep channel the river cuts on the east side where the rip drains, the
reef shelf on the west that makes the peak, and the notoriously steep shorebreak)
in analytic form. Drop a NOAA Hawaii coastal DEM into `bed()` for the true one.
"""
import argparse
import math
import urllib.request

import numpy as np
import torch
from kernels import get_kernel

BUOY = "https://www.ndbc.noaa.gov/data/realtime2/51201.spec"
G = 9.81
HDRY = 1e-3

# A representative Eddie-day sea state at 51201: the contest needs 20 ft Hawaiian
# faces, which shows up on the buoy as roughly 7.5 m at 17 s out of the NNW.
WINTER = dict(swh=7.5, swp=17.0, wwh=1.2, wwp=8.0, mwd=325.0)


def fetch_sea_state(timeout=20):
    """Latest swell + wind-wave partition from NDBC 51201."""
    with urllib.request.urlopen(BUOY, timeout=timeout) as r:
        rows = [l.split() for l in r.read().decode().splitlines() if not l.startswith("#")]
    for row in rows:                       # newest first; skip rows with missing data
        try:
            swh, swp, wwh, wwp = float(row[6]), float(row[7]), float(row[8]), float(row[9])
            mwd = float(row[14])
        except (ValueError, IndexError):
            continue
        return dict(swh=swh, swp=swp, wwh=wwh, wwp=wwp, mwd=mwd,
                    when=f"{row[0]}-{row[1]}-{row[2]} {row[3]}:{row[4]}Z")
    raise RuntimeError("no complete observation in 51201.spec")


def bed(X, Y, Lx, Ly):
    """Idealized Waimea Bay bed elevation (m, negative = below sea level)."""
    # cross-shore: deep outside, reef shelf, then a steep 1:10 beach face
    b = -30.0 + 22.0 * torch.sigmoid((X - 300.0) / 60.0)          # 30 m -> ~8 m shelf
    b = b + ((X - 460.0) / 10.0).clamp(min=0.0)                    # steep shorebreak
    # the west reef: a shallower lobe that trips the peak first
    reef = torch.exp(-(((X - 400.0) / 90.0) ** 2 + ((Y - 0.72 * Ly) / 110.0) ** 2))
    b = b + 3.2 * reef
    # the east channel: the river mouth stays deep, so the rip drains out here
    chan = torch.exp(-((Y - 0.18 * Ly) / 70.0) ** 2) * torch.sigmoid((X - 280.0) / 50.0)
    b = b - 5.0 * chan
    return b.clamp(max=4.0).contiguous()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--winter", action="store_true", help="representative Eddie-day swell")
    p.add_argument("--hs", type=float, help="override significant swell height (m)")
    p.add_argument("--tp", type=float, help="override swell period (s)")
    p.add_argument("--n", type=int, default=1024, help="cross-shore cells")
    p.add_argument("--frames", type=int, default=60)
    p.add_argument("--out", default="waimea")
    a = p.parse_args()

    if a.winter:
        ss = dict(WINTER, when="preset")
    else:
        ss = fetch_sea_state()
    if a.hs:
        ss["swh"] = a.hs
    if a.tp:
        ss["swp"] = a.tp
    print(f"51201 Waimea Bay [{ss['when']}]: swell {ss['swh']:.1f} m @ {ss['swp']:.1f} s, "
          f"windwave {ss['wwh']:.1f} m @ {ss['wwp']:.1f} s, MWD {ss['mwd']:.0f} deg")
    if not a.winter and ss["swh"] + ss["wwh"] < 2.0:
        print("  (the bay is flat today: North Shore is seasonal. try --winter)")

    surf = get_kernel("phanerozoic/surf", version=1, trust_remote_code=True)
    dev = "cuda"
    Nx = a.n
    Ny = int(a.n * 0.75)
    Lx, Ly = 800.0, 600.0
    dx, dy = Lx / Nx, Ly / Ny
    xs = torch.linspace(0, Lx, Nx, device=dev)
    ys = torch.linspace(0, Ly, Ny, device=dev)
    X, Y = torch.meshgrid(xs, ys, indexing="xy")
    b = bed(X, Y, Lx, Ly)

    h = (0.0 - b).clamp(min=0)
    hu = torch.zeros_like(h)
    hv = torch.zeros_like(h)
    foam = torch.zeros_like(h)

    h0 = 30.0
    dt = 0.30 * dx / math.sqrt(G * (h0 + ss["swh"]))
    cf = 0.004

    # obliquity: MWD is the direction waves come FROM, shore normal faces ~340 deg
    theta = math.radians(((ss["mwd"] - 340.0 + 180.0) % 360.0) - 180.0)
    a_sw, om_sw = ss["swh"] / 2.0, 2 * math.pi / max(ss["swp"], 1.0)
    a_ww, om_ww = ss["wwh"] / 2.0, 2 * math.pi / max(ss["wwp"], 1.0)
    k_sw = om_sw / math.sqrt(G * h0)
    k_ww = om_ww / math.sqrt(G * h0)
    xrel = 120.0
    alpha = ((xrel - X) / xrel).clamp(0, 1) ** 2         # sponge/wavemaker ramp

    def wavemaker(h, hu, hv, t):
        ky_s = k_sw * math.sin(theta)
        ky_w = k_ww * math.sin(theta)
        eta = (a_sw * torch.sin(om_sw * t - ky_s * Y)
               + a_ww * torch.sin(om_ww * t - ky_w * Y + 1.7))
        ht = (h0 + eta).clamp(min=0)
        u_in = eta * math.sqrt(G / h0)
        return ((1 - alpha) * h + alpha * ht,
                (1 - alpha) * hu + alpha * ht * u_in,
                (1 - alpha) * hv)

    def foam_step(foam, h, hu, hv):
        safe = h.clamp(min=0.05)
        u, v = hu / safe, hv / safe
        fr = torch.sqrt(u * u + v * v) / torch.sqrt(G * safe)
        wet = (h > 0.03).float()
        brk = ((fr - 0.75) / 0.4).clamp(0, 1) * wet
        return (foam * math.exp(-dt / 2.5) + dt * 6.0 * brk).clamp(0, 5)

    t = 0.0
    for _ in range(int(60.0 / dt)):                      # spin up a steady surf zone
        h, hu, hv = wavemaker(h, hu, hv, t)
        h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt, cf=cf)
        foam = foam_step(foam, h, hu, hv)
        t += dt

    sub = max(1, round(0.2 / dt))
    etaF, foamF = [], []
    for _ in range(a.frames):
        etaF.append(torch.where(h > HDRY, h + b, b).float().cpu().numpy())
        foamF.append(foam.float().cpu().numpy())
        for _ in range(sub):
            h, hu, hv = wavemaker(h, hu, hv, t)
            h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt, cf=cf)
            foam = foam_step(foam, h, hu, hv)
            t += dt

    et, fo = np.stack(etaF), np.stack(foamF)
    bn = b.float().cpu().numpy()
    np.save(f"{a.out}_eta.npy", et)
    np.save(f"{a.out}_foam.npy", fo)
    np.save(f"{a.out}_bed.npy", bn)

    print(f"finite={np.isfinite(et).all()}  eta[{et.min():.2f}, {et.max():.2f}] m")
    print("  cross-shore transect through the peak (y = 0.72 Ly):")
    j = int(0.72 * Ny)
    crest = et.max(axis=0)
    for xm in (150, 300, 380, 420, 460, 500, 540):
        i = min(int(xm / Lx * Nx), Nx - 1)
        d = -bn[j, i]
        print(f"    x={xm:3d} m  depth={d:6.1f} m  crest={crest[j, i]:5.2f} m  "
              f"foam={fo[:, j, i].max():4.2f}")


if __name__ == "__main__":
    main()