#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Apr 26 14:10:44 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import geoxarray
from misc.landmask import flagland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
dsname = 'Tas'
daname = 'Tas'
region = 'global'
sel_region = lambda x: x
season = 'annual'
years = slice('1870', '2020')
landonly = False
if 'land' in sys.argv: landonly = True
if 'NEUS' in sys.argv:
    region = 'NEUS'
    lons = slice(360-78, 360-70)
    lats = slice(36, 44)
    sel_region = lambda x: x.sel(lon=lons, lat=lats)
    region_long = f'{360-lons.start}-{360-lons.stop}W {lats.start}-{lats.stop}N'
    

title = f'{region} {season} {daname} {years.start}-{years.stop}'
if landonly: title += ' land'
if region not in ('global',): title = title.replace(f'{region}', f'{region} {region_long}') 
tag = title.replace(' ', '_')

if landonly:
    mask = lambda x: x.where(flagland(x)>0.5)
else:
    mask = lambda x: x

ofile = __file__.replace('.py', f'_{tag}.nc')
if os.path.exists(ofile):
    da = xr.open_dataarray(ofile).load()
    print('[loaded]:', ofile)
else:
    #no urban
    ifile = f'data_{dsname}_AM4urban_wasteCool_0urban_amip_10ens_1870-2020.nc'
    da0 = xr.open_dataarray(ifile).load() \
        .pipe(mask) \
        .pipe(sel_region) \
        .geo.fldmean() \
        .groupby('time.year').mean('time')
    #urban
    ifile = f'data_{dsname}_AM4urban_wasteCool_amip_10ens_1870-2020.nc'
    da = xr.open_dataarray(ifile).load() \
        .pipe(mask) \
        .pipe(sel_region) \
        .geo.fldmean() \
        .groupby('time.year').mean('time')
    
    da = xr.concat([da0, da], dim=pd.Index(['noUrban', 'urban'], name='case'))
    #savedata
    da.to_dataset(name=daname).to_netcdf(ofile)
    print('[saved]:', ofile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    plt.close('all')

    fig, ax = plt.subplots()
    convert_units = lambda x: x.where(x<100, other=x-273.15).assign_attrs(units='$^\circ$C')
    da.mean('en').pipe(convert_units).plot(hue='case')

    title += f' {da.en.size}ens'
    tag = title.replace(' ', '_')
    plt.title(title, loc='left')
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{tag}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    fig, ax = plt.subplots()
    lowpass = lambda x: x.filter.lowpass(1/11, dim='year', padtype='even')

    daa = (da.sel(case='urban') - da.sel(case='noUrban')).mean('en')
    daa.pipe(convert_units).plot()
    daa.pipe(lowpass).pipe(convert_units).plot(label='11-year lowpass')

    plt.legend()
    title += f' urban-noUrban'
    tag = title.replace(' ', '_')
    plt.title(title, loc='left')
    print(f'{daa.sel(year=slice(1901,2020)).mean().item() = }')

    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{tag}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    tt.check(f'**Done**')
    print()
    plt.show()