#!/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 = 'urbanfrac'
daname = 'urbanfrac'
region = 'global'
sel_region = lambda x: x
years = slice('1870', '2020')
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} {daname} {years.start}-{years.stop}'
if region not in ('global',): title = title.replace(f'{region}', f'{region} {region_long}') 
tag = title.replace(' ', '_')

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

ofile = __file__.replace('.py', f'_{tag}.nc')
if os.path.exists(ofile):
    da = xr.open_dataarray(ofile).load()
    print('[loaded]:', ofile)
else:
    #urban
    ifile = f'data_{dsname}_AM4urban_wasteCool_amip_1870-2020.nc'
    da = xr.open_dataarray(ifile).load() \
        .pipe(mask) \
        .pipe(sel_region) \
        .geo.fldmean()
    
    #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()
    da.plot()

    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)

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