# npgo_update.py -- project CMEMS sea level anomalies onto the NPGO pattern
#
# Minimal reproduction of the satellite-era NPGO update described at
# https://www.npgo.org/method.html. Needs: copernicusmarine, xarray, numpy,
# pandas, and a free Copernicus Marine account (run `copernicusmarine login`
# once). Place NPGO_EOF.nc and NPGO.txt (both from npgo.org) alongside.
#
# Cite Di Lorenzo et al. (2008), GRL, doi:10.1029/2007GL032838.

import copernicusmarine, numpy as np, pandas as pd, xarray as xr

BOX = dict(minimum_longitude=-180, maximum_longitude=-110,
           minimum_latitude=25,    maximum_latitude=62)

# 1. download DT monthly (1993-) and NRT daily (recent) SLA over the box
copernicusmarine.subset(
    dataset_id="cmems_obs-sl_glo_phy-ssh_my_allsat-l4-duacs-0.125deg_P1M-m",
    variables=["sla"], output_filename="sla_dt.nc", **BOX)
copernicusmarine.subset(
    dataset_id="cmems_obs-sl_glo_phy-ssh_nrt_allsat-l4-duacs-0.125deg_P1D",
    variables=["sla"], output_filename="sla_nrt.nc", **BOX)

# 2. monthly means; 3. splice NRT after the last DT month
dt  = xr.open_dataset("sla_dt.nc")["sla"].resample(time="1MS").mean()
nrt = xr.open_dataset("sla_nrt.nc")["sla"].resample(time="1MS").mean(skipna=False)
sla = xr.concat([dt, nrt.sel(time=nrt.time > dt.time[-1])], dim="time")

# 4. anomalies relative to the 1993-2004 monthly climatology
clim = sla.sel(time=slice("1993", "2004")).groupby("time.month").mean()
ssha = sla.groupby("time.month") - clim

# 5. regrid onto the EOF grid; 6. project
eof = xr.open_dataset("NPGO_EOF.nc")
pattern = eof["eof"] * eof["mask"]
ssha_i = ssha.interp(longitude=eof.lon, latitude=eof.lat)          # bilinear
valid = np.isfinite(pattern) & np.isfinite(ssha_i).all("time")      # common points
pc = (ssha_i.where(valid) * pattern.where(valid)).sum(("lon", "lat"))

# 7. normalize to the published index over the overlap
ref = np.loadtxt("NPGO.txt", comments="#")                          # year month value
ref_s = pd.Series(ref[:, 2], index=pd.PeriodIndex(
    [pd.Period(year=int(y), month=int(m), freq="M") for y, m in ref[:, :2]]))
pc_s = pd.Series(pc.values, index=pd.PeriodIndex(pc.time.values, freq="M"))
ov = pc_s.index.intersection(ref_s.index)
npgo_sat = (pc_s - pc_s[ov].mean()) / pc_s[ov].std() * ref_s[ov].std() + ref_s[ov].mean()

# 8. append the months beyond the published index
npgo = pd.concat([ref_s, npgo_sat[npgo_sat.index > ref_s.index[-1]]])
print(npgo.tail(6))