#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sat Feb 13 13:34:41 EST 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
import matplotlib.pyplot as plt
#more imports
import salem
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
dfsh= salem.read_shapefile('chn_admbnda_adm1_ocha/chn_admbnda_adm1_ocha.shp') 
dfsh['coords'] = dfsh['geometry'].apply(lambda x: x.representative_point().coords[:])
#dfsh['coords'] = dfsh['geometry'].apply(lambda x: x.centroid.coords[:])
dfsh['coords'] = [coords[0] for coords in dfsh['coords']]
dfsh['rank'] = dfsh.index
tag = []
for idx, row in dfsh.iterrows():
    tag.append(str(row['rank']) + ' ' + row['ADM1_EN'].split()[0]) 
dfsh['tag'] = tag
#clim data
da = xr.open_dataarray('era5.t2m.daily.China.subregion.2009-2014.nc')
dfsh['climdata'] = da.mean('time').values - 273.15
units, long_name = da.attrs['units'], da.attrs['long_name']
units = '$^\circ$C'
years = da.time.dt.year.values
title = f'{long_name}\n {years[0]}-{years[-1]} annual mean climatology [{units}]'
cmap = 'parula'
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #dfsh.plot(figsize=(10,8), column='ADM1_PCODE')
    dfsh.plot(figsize=(10,6), column='climdata', cmap=cmap, legend=True)
    for idx, row in dfsh.iterrows():
        plt.annotate(s=row['rank'], xy=row['coords'], fontsize='x-small',
                         horizontalalignment='center', va='center')

    ax = plt.gca()
    ax.axis('off')
    ax.set_title(title)
    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()