#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun 24 08:52:18 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'data/nino34.ctl.nyears1000.nc'
da = xr.open_dataarray(ifile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    nyears_rolling = 30
    da.drop('time').rolling(time=nyears_rolling*12+1, center=True).std().plot(ax=ax)
    #iyear = (year - 1)*12
    ax.axvspan((11-1)*12, (41-1)*12, color='k', alpha=1/8)
    ax.axvspan((301-1)*12, (331-1)*12, color='k', alpha=1/8)
    ax.axvspan((501-1)*12+1, (531-1)*12+1, color='k', alpha=1/8)
    da_ = da.drop('time').rolling(time=30*12+1, center=True).std()
    err = da_.std()
    da_ = da_.mean()
    ax.axhspan(da_-err, da_+err, color='k', alpha=1/8)
    ax.axhline(da.drop('time').rolling(time=30*12+1, center=True, min_periods=1).std().mean(), color='k', alpha=1/8, ls='--')
    
    ax.set_xticks(np.arange((100-1)*12, (1001-1)*12, 100*12))
    ax.set_xticklabels([f'{yyyy:04d}' for yyyy in range(100, 1001, 100)])
    ax.set_title(f'{nyears_rolling}-year running STD of Nino3.4 index')
    ax.set_ylabel('K')
    #ax.set_xlim(0, 1000*12)
    ax.set_xlim((300-1)*12, (400-1)*12)
    ax.set_xticks(np.arange((300-1)*12, (400-1)*12, 10*12))
    ax.set_xticklabels([f'{yyyy:04d}' for yyyy in range(300, 400, 10)])

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()