model = gcm1
periods = np.arange(1, 9)
longitudes = [-110]  # -140, -155]
tslice = slice(None, "1996-03-01")
subset_fields = [
    "Jq",
    "theta",
    "salt",
    "dens",
    "u",
    "v",
    "w",
    "KPP_diffusivity",
    "sst",
    "Ri",
]
anomalize_fields = ["u", "v", "w", "theta", "salt", "sst"]
anomalize_over = ["time"]

import xfilter

surf = (
    model.surface.sel(longitude=longitudes, method="nearest")
    .assign_coords(longitude=longitudes)
    .drop("depth")
)
tao = model.tao.sel(longitude=longitudes, latitude=0).chunk({"longitude": 1})
# tao["time"] = model.full.time

sst = surf.theta.sel(longitude=longitudes).chunk({"longitude": 1, "time": -1})
sstfilt = pump.calc.tiw_avg_filter_sst(sst, "bandpass")  # .persist()

# filter out daily cycle
subdaily_sst = xfilter.lowpass(
    surf.theta.chunk({"time": -1, "latitude": 10}),
    coord="time",
    freq=2,
    cycles_per="D",
).reindex(time=tao.time)
subdaily_sst.attrs["long_name"] = "2 day filtered SST"
subdaily_sst.attrs["units"] = "C"

# surf["theta"] = (
#    sst.pipe(xfilter.lowpass, coord="time", freq=1/7, cycles_per="D", num_discard=0, method="pad")
# )

Tx = model.surface.theta.differentiate("longitude")
Ty = model.surface.theta.differentiate("latitude")
gradT = (
    np.hypot(Tx, Ty)
    .sel(longitude=longitudes, method="nearest")
    .assign_coords(longitude=longitudes)
    .drop_vars("depth")
)

# with performance_report("profiles/phase-period.html"):
tiw_phase, period, tiw_ptp = dask.compute(
    pump.calc.get_tiw_phase_sst(
        sstfilt.chunk({"longitude": 1}), gradT.chunk({"longitude": 1, "time": -1})
    ),
)[0]

full = pump.utils.read_gcm1_zarr_subset(gcm1)

# needed for apply_ufunc later
valid_subset_fields = list(set(subset_fields).intersection(set(full.data_vars)))
full_subset = full[valid_subset_fields]

full_subset = pump.calc_reduced_shear(
    full_subset.sel(depth=slice(0, -200))
    .sel(longitude=longitudes)
    .sel(time=tslice)
    .chunk({"longitude": 1, "latitude": -1, "time": 36})
    .assign_coords(
        period=period,
        tiw_phase=tiw_phase,
    )
)
# full_subset["uz"] = full_subset.u.differentiate("depth")
# full_subset["vz"] = full_subset.v.differentiate("depth")
# full_subset["S2"] = full_subset.uz ** 2 + full_subset.vz ** 2
# full_subset["N2"] = -9.81/1025 * full_subset.dens.differentiate("depth")
# full_subset["Ri"] = full_subset.N2 / full_subset.S2
full_subset["sst"] = surf.theta.sel(
    time=full_subset.time
)  # full_subset.theta.isel(depth=0)
# if "zeta" in surf:
#    full["ssvor"] = surf.zeta

euc_max = tao.euc_max.reset_coords(drop=True)
mld = pump.calc.get_mld(full_subset.dens)
dcl_base = pump.calc.get_dcl_base_Ri(full_subset, mld, euc_max)


full_subset = full_subset.assign_coords(mld=mld, dcl_base=dcl_base, eucmax=tao.euc_max)

subset = (
    full_subset[["u", "v"]].sel(
        longitude=-110,
        latitude=slice(-3, 6),
        time=(full_subset.period == 4).squeeze().reset_coords(drop=True),
    )
).compute()

ssta4 = (
    gcm1.surface.theta.sel(time=full_subset.time)
    .sel(
        longitude=slice(-140, -100),
        time=(full_subset.period == 4).squeeze().reset_coords(drop=True),
    )
    .pipe(lambda x: x - x.mean(["longitude", "time"]))
    .reset_coords(drop=True)
)

ssta4.isel(time=slice(None, None, 20)).plot(
    x="longitude", y="latitude", col="time", col_wrap=3, robust=True, aspect=2.5
)

vort = xr.Dataset()
vort["x"] = ddy(ds.w) - ds.v.differentiate("depth")
vort["y"] = ds.u.differentiate("depth") - ddx(ds.w)
vort["z"] = ddx(ds.v) - ddy(ds.u)
vort["vx"] = ddx(ds.v)
vort["uy"] = ddy(ds.u)
vort = vort.isel(longitude=1).load()

subset = (
    vort.sel(depth=(-20, -40, -60), method="nearest")
    .sel(latitude=slice(-2, 5))
    .isel(latitude=slice(None, None, 5), time=slice(None, None, 5))
)

fg = quiver(
    subset,
    u="x",
    v="y",
    x="time",
    y="latitude",
    col="depth",
    scale=0.4,
)


uvsub = (
    ds.isel(longitude=1)[["u", "v"]]
    # .differentiate("depth")
    .sel(latitude=slice(-2, 5)).isel(
        latitude=slice(None, None, 6), time=slice(None, None, 6)
    )
).load()

for loc, ax in zip(fg.name_dicts.flat, fg.axes.flat):
    quiver(
        uvsub.sel(loc).load(),
        u="u",
        v="v",
        x="time",
        y="latitude",
        scale=10,
        ax=ax,
        color="r",
    )

subset = vort.sel(depth=(-20, -40, -60), method="nearest").sel(latitude=slice(-2, 5))

fg = subset.z.sel(depth=(-20, -40, -60), method="nearest").plot(
    x="time", col="depth", vmax=2.5e-5, cbar_kwargs={"label": "vertical vorticity"}
)
subset.time.attrs["name"] = ""

for loc, ax in zip(fg.name_dicts.flat, fg.axes.flat):
    quiver(
        subset.isel(latitude=slice(None, None, 5), time=slice(None, None, 5))
        .sel(loc)
        .load(),
        u="x",
        v="y",
        x="time",
        y="latitude",
        scale=0.3,
        ax=ax,
        color="k",
    )

fg.set_xlabels("")
fg.fig.suptitle(
    "Horizontal vortex vectors[$ω_xi + ω_yj$] over vertical vorticity [color]", y=1.03
)
plt.gcf().savefig("images/110-period-4-horizontal-vertical-vorticity-depth.png")

subset = vort.sel(latitude=slice(-2, 6), depth=slice(0, -60)).mean("depth")
f0 = f.reindex_like(subset.z)

subset.z.attrs["long_name"] = "vert vorticity"
((f0 + subset.vx - subset.uy)).plot(x="time", vmax=3e-5)
(
    ds.theta.isel(depth=0, longitude=1)
    .resample(time="D")
    .mean()
    .plot.contour(x="time", colors="k", levels=[23.8], zorder=2, add_labels=False)
)
quiver(
    subset.isel(latitude=slice(None, None, 4), time=slice(None, None, 4)),
    u="x",
    v="y",
    x="time",
    y="latitude",
    scale=0.3,
)

quiver(
    ds.sel(depth=-30, method="nearest").isel(
        longitude=1, latitude=slice(None, None, 4), time=slice(None, None, 6)
    ),
    u="u",
    v="v",
    x="time",
    y="latitude",
    color="darkgreen",
    scale=12,
)


plt.gcf().suptitle(
    "Horizontal vorticity vectors over ($f+v_x - u_y$) [color].\n Depth average to 60m",
    y=1.05,
)

# plt.gcf().savefig("images/110-period-4-horizontal-vertical-vorticity.png")

quiver(subset, x="time", y="latitude", u="u", v="v", scale=10)

subset = (
    ds.isel(longitude=1)
    .sel(
        latitude=slice(
            -2,
            6,
        ),
        depth=slice(-20, -60),
    )
    .isel(latitude=slice(0, -1, 2))
    .mean("depth")
    .resample(time="D")
    .mean()
)

t = (subset.time - subset.time[0]).values.astype("timedelta64[s]").astype("float32")

x = t * 1200e3 / 25 / 86400

y = subset.latitude  # np.linspace(-2, 6, 100)

start_points = [x[-4] * np.ones_like(y), y]
start_points = np.stack(start_points).T

plt.streamplot(
    x,
    y,
    subset.u.transpose().values,
    subset.v.transpose().values,
    integration_direction="forward",
    start_points=start_points,
)

plot_shear_terms(duzdt.sel(depth=[-25, -50, -75], method="nearest"))

plot_shear_terms(dvzdt.sel(depth=[-25, -50, -75], method="nearest"))

from datetime import timedelta as delta

import numpy as np
import xarray as xr

from parcels import (
    AdvectionRK4,
    FieldSet,
    JITParticle,
    ParticleSet,
    plotTrajectoriesFile,
)

variables = {"U": "u", "V": "v"}
dimensions = {"lon": "longitude", "lat": "latitude", "time": "time"}
fset = FieldSet.from_xarray_dataset(gcm1.surface, variables, dimensions)
# fset.add_periodic_halo(zonal=True, meridional=False, halosize=2)

fset

lons, lats = np.meshgrid(np.arange(-120.0, -105.0, 0.2), np.arange(-2, 8, 0.5))
pset = ParticleSet(fset, JITParticle, lon=lons, lat=lats)
ofile = pset.ParticleFile("particles.nc", outputdt=delta(hours=1))

kernel =

def beach(particle, fieldset, time):
    if particle.lon > -96:
        particle.lon = -96

gcm1.surface

# %%time

# Advect ParticleSet with RK4 and periodic Boundary conditions
pset.execute(
    pset.Kernel(AdvectionRK4) + pset.Kernel(beach),
    runtime=delta(days=60),
    dt=delta(minutes=30),
    output_file=ofile,
)

ds = xr.open_dataset("particles.nc").set_coords(["time"])
ds["time"] = ds.time.isel(traj=0, drop=True)
ds = ds.rename({"obs": "time"})
ds["trajectory"] = ds.trajectory.isel(time=0)
ds = ds.rename({"trajectory": "traj"})
ds = ds.set_coords("traj")
ds

ds.sel(traj=slice(None, None, 10)).plot.scatter(x="lon", y="lat", s=0.05)

dirname = "/glade/campaign/cgd/oce/people/dcherian/TPOS_MITgcm_1_hb/HOLD"
metrics = pump.model.read_metrics(f"{dirname}/../")
ds = xr.open_mfdataset(f"{dirname}/File_*.nc", parallel=True)
# RAC = xr.align(metrics.RAC, ds, join="right")[0].chunk(
#    {k: v for k, v in ds.chunks.items() if k in ["latitude", "longitude"]}
# )
# ds["Jq"] = 1035 * 3994 * (ds.DFrI_TH + ds.KPPg_TH) / RAC
ds