#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Feb 25 11:57:10 AM EST 2026
# https://ajdawson.github.io/eofs/examples/elnino_xarray.html
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
from misc import get_kws_from_argv
from eofs.xarray import Eof
#from eofs.examples import example_data_path
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
imode = get_kws_from_argv('imode', 1); imode = int(imode)
year0 = get_kws_from_argv('year0', '1981')
#sst data
filename = '/projects/w/wenchang/data/ersst/v5/ersst5.185401-202405.nc'
da  = xr.open_dataset(filename)['sst']
#sst = da.sel(time=slice(f'{year0}-01', '2023-12')).resample(time='AS').mean()
sst = da.sel(time=slice(f'{year0}-01', '2023-12')) 
sst = sst.isel(time=sst.time.dt.month==12).groupby('time.year').mean('time').rename(year='time')
#sst = da.sel(time=slice(f'{year0}-06', None)).resample(time='AS-Jun').mean()
#da = da.sel(time=slice(f'{year0}-06', '2023-12'))
#sst = da.groupby('time.month') - da.groupby('time.month').mean('time')
#print(sst); sys.exit()

# Create an EOF solver to do the EOF analysis. Square-root of cosine of
# latitude weights are applied before the computation of EOFs.
coslat = np.cos(np.deg2rad(sst.coords['lat'].values))
wgts = np.sqrt(coslat)[..., np.newaxis]
solver = Eof(sst, weights=wgts)

#EOF calculations
N = 10
eofs = solver.eofs(eofscaling=2, neofs=N)
pcs = solver.pcs(pcscaling=1, npcs=N)
#print(eofs, pcs); sys.exit()
vf = solver.varianceFraction(neigs=N)
vf_error = solver.northTest(vfscaled=True, neigs=N)
total_variance = solver.totalAnomalyVariance()
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #import cartopy.crs as ccrs
    #import cartopy.feature as cfeature
    from geoplots import mapplot
    fig,ax = plt.subplots()
    plt.errorbar(vf.mode+1, vf, yerr=vf_error, capsize=5)
    ax.set_ylabel('variance fraction')
    ax.set_xlabel('mode')
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__varianceFraction.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    fig,axes = plt.subplots(2, 1, figsize=(6,6)) 
    ax = axes[0]
    eofs.isel(mode=imode-1).plot.contourf(ax=ax, levels=21)
    # EOF
    #ax.add_feature(cfeature.LAND, facecolor='w', edgecolor='k')
    #cb = plt.colorbar(fill, orientation='horizontal')
    #cb.set_label('correlation coefficient', fontsize=12)
    plt.sca(ax)
    mapplot(fill_continents=True)
    ax.set_title(f'EOF{imode}')
    # Plot the leading PC time series.
    ax = axes[1]
    #pc1[:, 0].plot(color='b', linewidth=2)
    #pc1[:, 0].plot()
    pcs.isel(mode=imode-1).plot(ax=ax, marker='o', fillstyle='none')
    ax.axhline(0, color='gray', ls='--')
    ax.set_xlabel('Year')
    ax.set_ylabel('Normalized Units')
    ax.set_title(f'PC{imode} Time Series')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__EOF{imode}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()