import sys, os, pickle, copy
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib import ticker
import cartopy.crs as ccrs
import cartopy.feature as cf
from datetime import datetime
from dateutil.relativedelta import relativedelta
from utils import lm_ID, wrk_dir, out_dir, addCbar, borderMask
from utils_cities import cities
plt_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']

imgsDir = './imgs/paper/'

lcz_color = [ '#910613', '#D9081C', '#FF0A22', '#C54F1E', '#FF6628', '#FF985E', '#FDED3F', '#BBBBBB', '#FFCBAB', '#565656', '#006A18', '#00A926', '#628432', '#B5DA7F', '#000000', '#FCF7B1', '#656BFA']


#===============================================================================
#
ini_date = datetime(2017, 1, 1)
end_date = datetime(2023, 1, 1)
years = list(np.arange(ini_date.year,end_date.year))
n_years = len(years)

extpar = xr.open_dataset('/project/pr133/jcanton/lmp_template.jcanton/lib/extpar_CH_1100m_lcz_0_default_fixDC_True.nc')
lcz_areas = xr.open_dataset('/store/g142/jcanton/repos/WUDAPT-to-COSMO/data_CH/extpar_CH_1100m_lcz_areas.nc')
extpar = extpar.merge(lcz_areas)
cell_dx   = 1100  # [m]
cell_area = cell_dx**2 # [m2]
# Remember that the grid is larger
xtpr_skp = 10

#plt.figure(1); plt.clf(); plt.show(block=False)
#plt.pcolormesh(extpar.lon, extpar.lat, extpar['LCZ_11'], shading='nearest')
#plt.draw()

with open('../WUDAPT-to-COSMO/tables/LCZ_UCP_default.csv', 'r') as ifile:
    ucpTable = np.loadtxt(ifile, delimiter=';', skiprows=1)
    # class; ISA; FR_PAVED; URBAN; URB_BLDFR; URB_BLDH; URB_H2W;URB_RfALB;URB_WaALB;URB_RdALB;URB_RfEMI;URB_WaEMI;URB_RdEMI;URB_RfHCAP;URB_WaHCAP;URB_RdHCAP;URB_RfHCON;URB_WaHCON;URB_RdHCON;AHF

#===============================================================================
# Print city tables
#
grid_lon = extpar.lon.values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
grid_lat = extpar.lat.values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
AHF      = extpar['AHF'     ].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
FR_PAVED = extpar['FR_PAVED'].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
HSURF    = extpar['HSURF'   ].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]

# urb_vnames = ['FR_PAVED', 'URB_FR_BLD', 'URB_H_BLD', 'URB_H2W', 'URB_ALB_SO', 'URB_ALB_TH', 'URB_EMIS', 'URB_HCON', 'URB_HCAP']
# for vname in urb_vnames:
#     URB = extpar[vname].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
#     I, J = cities[0]['cityIJ']
#     zh = URB[J, I]
#     I, J = cities[1]['cityIJ']
#     ge = URB[J, I]
#     print('{:10s} {:10.2e} {:10.2e} {:10.2e} {:10.2f}'.format(vname, zh, ge, zh-ge, (zh-ge)/zh*100))

for city in cities:

    I, J = city['cityIJ']
    hC = HSURF[J, I]
    I, J = city['ruralIJ']
    hR = HSURF[J, I]

    print('{:10s}: {:10.2f} {:10.2f} {:10.2f}'.format(city['name'], hC, hR, hC-hR))

print('City      :       Area | Paved area |   LCZ area |        AHF |  AHF/Parea |   mean UHI |    max UHI | timmax UHI | ')
for city in cities:

    # mask
    mask = city['mask']
    c_area = np.sum(mask)*cell_area

    # AHF and area paved
    c_ahf = np.sum(AHF[mask]*cell_area)
    c_fr_paved = FR_PAVED[mask]
    c_area_paved = np.sum(c_fr_paved*cell_area)

    # comparison between lcz-covered area and fr_paved area
    c_fr_lcz = 0. * AHF[mask]
    for lcz_nr in range(1,11):
        lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
        c_fr_lcz += lcz[mask]
    c_area_lcz = np.sum(c_fr_lcz*cell_area)

    # mean and max UHI
    a_times   = np.array(city['times'])
    a_rural   = np.array(city['t2m_rural'])
    a_cities  = np.array(city['t2m_city'])
    a_uhis    = a_cities - a_rural
    c_uhi_avg = np.mean(a_uhis)
    c_uhi_max = np.max(a_uhis)
    c_uhi_max_time = a_times[np.argmax(a_uhis)]
    c_uhi_min = np.min(a_uhis)
    #
    # max multiply observed UHI
    nobs = 2
    bins = np.arange(np.min(a_uhis), np.max(a_uhis), 0.1)
    counts, bins = np.histogram(a_uhis, bins=bins)
    idx = np.where(counts>0)[0][-1]
    c_uhi_max2 = (bins[idx] + bins[idx+1])/2
    #
    # mean max day-/night-time UHI (Zhao 2014)
    a_asod = np.array(city['avg_asod_s'])
    daytime_uhi   = np.where(a_asod < 1e-1, 0, a_uhis)
    nighttime_uhi = np.where(a_asod < 1e-1, a_uhis, 0)
    day_uhis   = []
    night_uhis = []
    p0 = p1 = 0
    while p1 < len(a_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(a_times >= a_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(a_times)+1
        day_uhis  .append( np.max(  daytime_uhi[p0:p1]) )
        night_uhis.append( np.max(nighttime_uhi[p0:p1]) )
    day_uhis   = np.array(day_uhis)
    night_uhis = np.array(night_uhis)

    print('{0:10s}: {1:10.2f} | {2:10.2f} | {3:10.2f} | {4:10.2f} | {5:10.2f} | {6:10.2f} | {7:10.2f} | {8:s} |'.format( \
            city['name'], c_area/(1e6), c_area_paved/(1e6), c_area_lcz/(1e6), c_ahf/1000, c_ahf/c_area_paved, c_uhi_avg, c_uhi_max2, c_uhi_max_time.strftime('%Y%m%d %H%M') \
            ))
    city['ahf'       ] = c_ahf
    city['fr_paved'  ] = c_fr_paved
    city['area_paved'] = c_area_paved
    city['fr_lcz'    ] = c_fr_lcz
    city['area_lcz'  ] = c_area_lcz
    city['ahf_sqm'   ] = c_ahf / c_area_paved
    #
    city['avg_uhi'   ] = c_uhi_avg
    city['max_uhi'   ] = c_uhi_max
    city['min_uhi'   ] = c_uhi_min
    city['max_uhi2'  ] = c_uhi_max2
    #
    city['avg_max_day_uhi'  ] = np.mean(day_uhis)
    city['avg_max_night_uhi'] = np.mean(night_uhis)

#-------------------------------------------------------------------------------
# Figure
# Dependence of UHI on parameters
#
#labels = ['ahf', 'area_paved', 'ahf_sqm', 'population', 'sim_city_nights', 'avg_asod_s', 'avg_clct', 'avg_winM', 'avg_prec', 'sky-view', 'H2W']
vnames = [ 'ahf',  'area_paved',          'ahf_sqm',  'population',          'sky-view',               'H2W',  'avg_asod_s',  'avg_clct',  'avg_prec']
labels = [r'AHF', r'$\log(\mathrm{S})$', r'AHF/S',   r'$\log(\mathrm{P})$', r'$\psi_{\mathrm{sky}}$', r'H/W', r'ASDSWR',     r'CLCT',     r'precip.' ]

fig = plt.figure(313, figsize=(14,6)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(4, len(vnames), num=fig.number, sharex='col', sharey='row')

lcz_prcn = np.zeros((10,len(cities)))

for ix, vname in enumerate(vnames):
    for icity, city in enumerate(cities):

        if vname in ['sim_city_nights', 'avg_asob_s', 'avg_asod_s', 'avg_clct', 'avg_winM']:
            ax[0, ix].plot(np.mean(city[vname]), city['avg_uhi' ], 'o', label=city['name'])
            ax[1, ix].plot(np.mean(city[vname]), city['max_uhi2'], 'o', label=city['name'])
            ax[2, ix].plot(np.mean(city[vname]), city['avg_max_day_uhi'  ], 'o', label=city['name'])
            ax[3, ix].plot(np.mean(city[vname]), city['avg_max_night_uhi'], 'o', label=city['name'])
        elif vname in ['avg_prec']:
            ax[0, ix].plot(365*np.mean(city[vname]), city['avg_uhi' ], 'o', label=city['name'])
            ax[1, ix].plot(365*np.mean(city[vname]), city['max_uhi2'], 'o', label=city['name'])
            ax[2, ix].plot(365*np.mean(city[vname]), city['avg_max_day_uhi'  ], 'o', label=city['name'])
            ax[3, ix].plot(365*np.mean(city[vname]), city['avg_max_night_uhi'], 'o', label=city['name'])
        elif vname in ['area_paved', 'area_lcz', 'population']:
            ax[0, ix].plot(np.log(city[vname]), city['avg_uhi' ], 'o', label=city['name'])
            ax[1, ix].plot(np.log(city[vname]), city['max_uhi2'], 'o', label=city['name'])
            ax[2, ix].plot(np.log(city[vname]), city['avg_max_day_uhi'  ], 'o', label=city['name'])
            ax[3, ix].plot(np.log(city[vname]), city['avg_max_night_uhi'], 'o', label=city['name'])
        elif vname in ['sky-view', 'H2W']:
            # get the percentage of lcz classes for the city centre
            for lcz_nr in range(1,11):
                lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
                I, J = city['cityIJ']
                lcz_prcn[lcz_nr-1,icity] = lcz[J, I]
            #
            svf = [0.30, 0.45, 0.40, 0.60, 0.65, 0.75, 0.35, 0.70, 0.80, 0.75]
            h2w = [2.5 , 1.25, 1.25, 1.  , 0.5 , 0.5 , 1.5 , 0.2 , 0.15, 0.35]
            if vname == 'sky-view':
                avg_vals = svf
            elif vname == 'H2W':
                avg_vals = h2w
            ax[0, ix].plot(np.mean(lcz_prcn[:icity]*avg_vals), city['avg_uhi' ], 'o', label=city['name'])
            ax[1, ix].plot(np.mean(lcz_prcn[:icity]*avg_vals), city['max_uhi2'], 'o', label=city['name'])
            ax[2, ix].plot(np.mean(lcz_prcn[:icity]*avg_vals), city['avg_max_day_uhi'  ], 'o', label=city['name'])
            ax[3, ix].plot(np.mean(lcz_prcn[:icity]*avg_vals), city['avg_max_night_uhi'], 'o', label=city['name'])
        else:
            ax[0, ix].plot(city[vname], city['avg_uhi' ], 'o', label=city['name'])
            ax[1, ix].plot(city[vname], city['max_uhi2'], 'o', label=city['name'])
            ax[2, ix].plot(city[vname], city['avg_max_day_uhi'  ], 'o', label=city['name'])
            ax[3, ix].plot(city[vname], city['avg_max_night_uhi'], 'o', label=city['name'])

        ax[-1,ix].set_xlabel(labels[ix])

ax[0,0].set_ylabel(r'<UHI>')
ax[1,0].set_ylabel(r'max(UHI)')
ax[2,0].set_ylabel(r'<max(UHI$_{\mathrm{D}}$)>')
ax[3,0].set_ylabel(r'<max(UHI$_{\mathrm{N}}$)>')

for iax in range(len(ax)):
    box = ax[0, iax].get_position()
    ax[0, iax].set_position([box.x0, box.y0, box.width * 0.9, box.height])
    box = ax[1, iax].get_position()
    ax[1, iax].set_position([box.x0, box.y0, box.width * 0.9, box.height])

ax[1,-1].legend(loc='center left', bbox_to_anchor=(1, -0.1), fontsize='small')

plt.draw()
plt.savefig(imgsDir+'uhi_vs_params2.png')
plt.savefig(imgsDir+'uhi_vs_params2.pdf')

#-------------------------------------------------------------------------------
# LCZ table
print('LCZ & ', end='')
for lcz_nr in range(1,11):
    print('\\multicolumn{{2}}{{c}}{{{0:2d}}} & '.format(lcz_nr), end='')
print('\\\\')
print('City       &', end='')
for lcz_nr in range(1,11):
    print('     Area &  \\% &'.format(lcz_nr), end='')
print('\\\\')
for city in cities:

    mask = city['mask']

    print('{0:10s} &'.format(city['name']), end='')

    city['lcz_area'] = np.zeros(10)
    city['lcz_prcn'] = np.zeros(10)

    for lcz_nr in range(1,11):

        lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]

        lcz_area = np.sum(lcz[mask]*cell_area)
        lcz_prcn = lcz_area / city['area_lcz'] * 100

        print('{0:7.2f} & \\cellcolor{{green!{2:d}}}{{{1:4.1f}}} &'.format( \
                lcz_area/(1e6), lcz_prcn,int(lcz_prcn) \
                ), end='')
        city['lcz_area'][lcz_nr-1] = lcz_area
        city['lcz_prcn'][lcz_nr-1] = lcz_prcn

    print('\\\\')


# reference cells composition
print('City       : Reference rural cell composition')
for city in cities:

    print('{0:10s} :'.format(city['name']), end='')

    I, J = city['ruralIJ']

    for lcz_nr in range(1,11):
        lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
        print(' {0:4.1f}'.format(100*lcz[J, I]), end='')
    print()

print('City       : city centre composition')
for city in cities:

    print('{0:10s} :'.format(city['name']), end='')

    I, J = city['cityIJ']

    for lcz_nr in range(1,11):
        lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]
        print(' {0:4.1f}'.format(100*lcz[J, I]), end='')
    print()

#===============================================================================
# Try to put in a figure
#
fig = plt.figure(1, figsize=(10,6)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(3, 1, num=fig.number, sharex=True, sharey=True)
width=0.5
city_names = [city['name'] for city in cities]
bottom = np.zeros((3,len(cities)))
lcz_prcn = np.zeros((3,len(cities)))
for lcz_nr in range(1,11):

    lcz = extpar['LCZ_{0:02d}'.format(lcz_nr)].values[xtpr_skp:-xtpr_skp,xtpr_skp:-xtpr_skp]

    for icity, city in enumerate(cities):
        lcz_prcn[0,icity] = city['lcz_prcn'][lcz_nr-1]
        I, J = city['cityIJ']
        lcz_prcn[1,icity] = 100*lcz[J, I]
        I, J = city['ruralIJ']
        lcz_prcn[2,icity] = 100*lcz[J, I]

    ax[0].bar(city_names, lcz_prcn[0,:], width=width, bottom=bottom[0,:], color=lcz_color[lcz_nr-1], label='LCZ {0:2d}'.format(lcz_nr))
    ax[1].bar(city_names, lcz_prcn[1,:], width=width, bottom=bottom[1,:], color=lcz_color[lcz_nr-1], label='LCZ {0:2d}'.format(lcz_nr))
    ax[2].bar(city_names, lcz_prcn[2,:], width=width, bottom=bottom[2,:], color=lcz_color[lcz_nr-1], label='LCZ {0:2d}'.format(lcz_nr))

    bottom += lcz_prcn

ax[0].set_ylabel('City')
ax[1].set_ylabel('Centre')
ax[2].set_ylabel('Rural')
ax[2].set_xticklabels(city_names, rotation=45, ha='right')
ax[0].grid(which='major', axis='y')
ax[1].grid(which='major', axis='y')
ax[2].grid(which='major', axis='y')
ax[0].set_axisbelow(True)
ax[1].set_axisbelow(True)
ax[2].set_axisbelow(True)
# Shrink the axes
for iax in range(len(ax)):
    box = ax[iax].get_position()
    ax[iax].set_position([box.x0, box.y0, box.width * 0.9, box.height])
# Put a legend to the right of the current axis
ax[1].legend(loc='center left', bbox_to_anchor=(1, 0.5), fontsize='small')
plt.draw()
plt.savefig(imgsDir+'cities_composition.pdf')
plt.savefig(imgsDir+'cities_composition.png')

#===============================================================================
# UHI spatial maps
#
months = [6, 7, 8]
# load and average
for iyear, year in enumerate(years):
    for imonth, month in enumerate(months):

        sim_date = datetime(year, month, 1)
        sim_dir = sim_date.strftime('%y%m0100') + lm_ID
        fpath = os.path.join(wrk_dir, sim_dir, out_dir, 'avg_T_2M.nc')
        T2M = xr.load_dataset(fpath)['T_2M']
        lon = T2M.lon.values
        lat = T2M.lat.values
        if iyear==0 and imonth==0:
            t2m = T2M.values[0,:,:]
        else:
            t2m += T2M.values[0,:,:]

t2m = t2m / (len(months)*len(years))


fig = plt.figure(2, figsize=(11.5,6.5)); plt.clf(); plt.show(block=False)
for icity, city in enumerate(cities):

    ax = fig.add_subplot(2, 5, icity+1, projection=ccrs.PlateCarree())

    #bndr_mask = city['mask']
    #bndr_mask = borderMask(city['mask'], nr_layers=1, diags=False)
    bndr_mask = borderMask(city['mask'], nr_layers=20, diags=True)

    lonMin = grid_lon[bndr_mask].min()
    lonMax = grid_lon[bndr_mask].max()
    latMin = grid_lat[bndr_mask].min()
    latMax = grid_lat[bndr_mask].max()

    uhi = copy.copy(t2m)
    #uhi[~bndr_mask] = 1e9
    #refT = uhi.min()
    #stationIJ = city['stationIJ']
    #refT = uhi[stationIJ[1], stationIJ[0]]
    ruralIJ = city['ruralIJ']
    refT = uhi[ruralIJ[1], ruralIJ[0]]
    uhi = uhi - refT
    uhi[~bndr_mask] = np.nan

    varMin = uhi[bndr_mask].min()
    varMax = uhi[bndr_mask].max()

    #im = ax.pcolormesh(lon, lat, uhi, vmin=varMin, vmax=varMax, shading='auto', cmap='cool')
    #im = ax.contourf(lon, lat, uhi, levels=np.linspace(0,varMax,10), cmap='cool')

    nlevels = 40
    norm = colors.TwoSlopeNorm(vmin=varMin, vcenter=0, vmax=varMax)
    levels = np.r_[np.linspace(varMin,0,nlevels//2)[0:-1], np.linspace(0,varMax,nlevels//2)]

    im = ax.contourf(grid_lon, grid_lat, uhi, levels=levels, cmap='seismic', norm=norm)

    #ax.set_title(city['name'])
    #if (icity==7):
    #    addCbar(fig, ax, im, label='UHI [K]', lonMM=[lonMin,lonMax], latMM=[latMin,latMax], extend='max', borders=False)
    #else:
    #    ax.axis('equal')
    #    ax.set_xlim([lonMin, lonMax])
    #    ax.set_ylim([latMin, latMax])
    #    ax.axis('off')
    #addCbar(fig, ax, im, label=city['name'] + ' UHI [K]', lonMM=[lonMin,lonMax], latMM=[latMin,latMax], borders=False)
    addCbar(fig, ax, im, label=city['name'] + ' UHI [K]', lonMM=[lonMin,lonMax], latMM=[latMin,latMax], borders=False, ticks=[levels[i] for i in [0,nlevels//2-1,nlevels-2]], dy=3)

    # save to file
    if 1 == 1:
        uhi[~bndr_mask] = -9999
        #ds = xr.Dataset(
        #        attrs={'title': 'UHI {0:s}'.format(city['name'])},
        #        data_vars={
        #            'UHI': xr.DataArray(
        #                data   = uhi,
        #                dims   = ['lat', 'lon'],
        #                coords = {'latitude': (['lat', 'lon'], lat), 'longitude': (['lat', 'lon'], lon)},
        #                attrs  = {
        #                    'standard_name' : 'urban_heat_island',
        #                    'long_name'     : 'Urban heat island',
        #                    '_FillValue'    : 0,
        #                    'missing_value' : -9999,
        #                    }
        #                )
        #            }
        #        )

        # write to file
        ds = xr.load_dataset(fpath)
        ds.drop(labels=['T_2M'])
        ds = ds.assign(variables = {'UHI': ( ('rlat','rlon'), uhi ) })
        ds.to_netcdf('./pp_data/UHI_{0:s}.nc'.format(city['name'].replace('.', '')))

plt.draw()
#plt.savefig(imgsDir+'uhi_map_{0:4d}m{1:s}.pdf'.format(year,''.join([str(month) for month in months])))

#===============================================================================
# Figure
# Time series of daily max UHI
#
fig = plt.figure(3, figsize=(12,8)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(2, 1, num=fig.number, sharex=True, sharey=False)
for icity, city in enumerate(cities):


    all_times  = np.array(city['times'])
    all_rural  = np.array(city['t2m_rural'])
    all_cities = np.array(city['t2m_city'])
    days = []
    uhis = []

    p0 = p1 = 0
    while p1 < len(all_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(all_times >= all_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(all_times)+1

        days.append( all_times[p0] )
        uhis.append( np.max( all_cities[p0:p1] - all_rural[p0:p1] ) )
        #uhis.append( np.mean( all_cities[p0:p1] - all_rural[p0:p1] ) )

    days = np.array(days)
    uhis = np.array(uhis)
    #ax[0].plot(days, uhis, label=city['name'])

    # monthly mean value
    mo_days = []
    mo_uhis = []
    mo_stds = []
    p0 = p1 = 0
    while p1 < len(days):
        p0 = p1
        try:
            p1 = np.argwhere(days >= days[p0]+relativedelta(months=1))[0][0]
        except:
            p1 = len(days)+1

        mo_days.append( days[p0] )
        mo_uhis.append( np.mean( uhis[p0:p1] ) )
        mo_stds.append( np.std(  uhis[p0:p1] ) )

    mo_uhis = np.array(mo_uhis)
    mo_stds = np.array(mo_stds)

    ax[0].plot(mo_days, mo_uhis, label=city['name'])
    # too many coloured areas
    #ax[0].fill_between(mo_days, mo_uhis-mo_stds, mo_uhis+mo_stds, alpha=0.5)

    # looks terrible
    # # average over years
    # ma_days = mo_days[0:12]
    # ma_uhis = np.zeros(12)
    # for yy in range(5):
    #     ma_uhis += mo_uhis[12*yy: 12*(yy+1)]
    # ma_uhis /= 5
    # ax[0].plot(ma_days, ma_uhis, label=city['name'])


    # yearly mean value
    ye_days = []
    ye_uhis = []
    p0 = p1 = 0
    while p1 < len(days):
        p0 = p1
        try:
            p1 = np.argwhere(days >= days[p0]+relativedelta(years=1))[0][0]
        except:
            p1 = len(days)+1

        ye_days.append( days[p0] )
        ye_uhis.append( np.mean( uhis[p0:p1] ) )

    ax[1].plot(ye_days, ye_uhis, label=city['name'])

ax[0].grid(which='major', axis='x')
ax[1].grid(which='major', axis='x')
ax[-1].legend(fontsize='small', loc=7)
ax[-1].set_xlim([datetime(2017,1,1), datetime(2022,9,1)])
ax[0].set_ylabel('Monthly-averaged max UHI')
ax[1].set_ylabel('Yearly-averaged max UHI')
plt.draw()
plt.savefig(imgsDir+'uhi_timeSeries_atCities')
plt.savefig(imgsDir+'uhi_timeSeries_atCities.pdf')

#===============================================================================
# Figure
# same as above but averaging the months
#
fig = plt.figure(31, figsize=(12,6.4)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(2, 1, num=fig.number, sharex=False, sharey=False)
mo_all = np.zeros(12)
ye_all = np.zeros(len(years))
for icity, city in enumerate(cities):

    all_times  = np.array(city['times'])
    all_rural  = np.array(city['t2m_rural'])
    all_cities = np.array(city['t2m_city'])
    days = []
    uhis = []

    p0 = p1 = 0
    while p1 < len(all_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(all_times >= all_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(all_times)+1

        days.append( all_times[p0] )
        #uhis.append( np.max( all_cities[p0:p1] - all_rural[p0:p1] ) )
        uhis.append( np.mean( all_cities[p0:p1] - all_rural[p0:p1] ) )

    days = np.array(days)
    uhis = np.array(uhis)
    #ax[0].plot(days, uhis, label=city['name'])

    # monthly mean value
    mo_months = np.arange(12)
    mo_uhis   = np.zeros(12)
    mo_stds   = np.zeros(12)
    mo_counts = np.zeros(12)
    p0 = p1 = 0
    while p1 < len(days):
        p0 = p1
        try:
            p1 = np.argwhere(days >= days[p0]+relativedelta(months=1))[0][0]
        except:
            p1 = len(days)+1

        month = days[p0].month

        mo_uhis[month-1] += np.mean( uhis[p0:p1] )
        mo_stds[month-1] += np.std(  uhis[p0:p1] )
        mo_counts[month-1] += 1

    mo_uhis = mo_uhis / mo_counts
    mo_stds = mo_stds / mo_counts

    mo_all += mo_uhis

    ax[0].plot(mo_months+1, mo_uhis,                     '-o',            color=plt_colors[icity], label=city['name'])
    ax[0].plot(mo_months+1, np.mean(mo_uhis)*np.ones(12), '-', alpha=0.5, color=plt_colors[icity], label='_nolegend')
    # too many coloured areas
    #ax[0].fill_between(mo_days, mo_uhis-mo_stds, mo_uhis+mo_stds, alpha=0.5)

    # yearly mean value
    ye_days = []
    ye_uhis = []
    p0 = p1 = 0
    while p1 < len(days):
        p0 = p1
        try:
            p1 = np.argwhere(days >= days[p0]+relativedelta(years=1))[0][0]
        except:
            p1 = len(days)+1

        ye_days.append( days[p0] )
        ye_uhis.append( np.mean( uhis[p0:p1] ) )

    ye_all += ye_uhis

    ax[1].plot(ye_days, ye_uhis, '-o', label=city['name'])

mo_all = mo_all / len(cities)
ye_all = ye_all / len(cities)
ax[0].plot(mo_months+1,         mo_all,                 '-d', color='black',            label='_nolegend_')
ax[0].plot(mo_months+1, np.mean(mo_all)*np.ones(12),     '-', color='black', alpha=0.5, label='_nolegend_')
ax[1].plot(ye_days,         ye_all,                     '-d', color='black',            label='_nolegend_')
ax[1].plot(ye_days, np.mean(ye_all)*np.ones(len(years)), '-', color='black', alpha=0.5, label='_nolegend_')

ax[0].grid(which='major', axis='x', alpha=0.3)
ax[1].grid(which='major', axis='x', alpha=0.3)
ax[1].legend(fontsize='small', loc=7)
ax[0].set_xlim([1-0.1, 12+0.1])
ax[0].set_xticks(np.arange(1,13))
ax[0].set_xticklabels(['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'])
ax[1].set_xlim([datetime(2016,12,15), datetime(2022,11,1)])
ax[0].set_ylabel(r'<UHI>$_{\mathrm{month}}$')
ax[1].set_ylabel(r'<UHI>$_{\mathrm{year}}$')
ax[0].text(-.08, .95, '(a)', transform=ax[0].transAxes)
ax[1].text(-.08, .95, '(b)', transform=ax[1].transAxes)
plt.draw()
plt.savefig(imgsDir+'uhi_timeSeries_atCities_b.png')
plt.savefig(imgsDir+'uhi_timeSeries_atCities_b.pdf')

#===============================================================================
# Figure
# histograms
#
fig = plt.figure(32, figsize=(9.68,3.00)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(1, 2, num=fig.number, sharex=False, sharey=False)
for icity, city in enumerate(cities):

    all_times  = np.array(city['times'])
    all_rural  = np.array(city['t2m_rural'])
    all_cities = np.array(city['t2m_city'])
    all_uhis   = all_cities - all_rural

    days = []
    avg_uhis = []
    max_uhis = []
    max_time = []

    p0 = p1 = 0
    while p1 < len(all_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(all_times >= all_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(all_times)+1

        days.append( all_times[p0] )
        avg_uhis.append( np.mean( all_uhis[p0:p1] ) )
        max_uhis.append( np.max(  all_uhis[p0:p1] ) )
        max_time.append( all_times[p0+np.argmax(  all_uhis[p0:p1] )].hour )

    days = np.array(days)
    ndays = len(days)
    avg_uhis = np.array(avg_uhis)
    max_uhis = np.array(max_uhis)
    max_time = np.array(max_time)

    #ax[0].hist(all_uhis, bins=50, density=False, histtype='step', log=False, label=city['name'])
    #ax[0].hist(avg_uhis, bins=50, density=False, histtype='stepfilled', alpha=0.4, log=False, label=city['name'])
    #ax[1].hist(max_uhis, bins=50, density=False, histtype='stepfilled', alpha=0.4, log=False, label=city['name'])

    avg_counts, avg_bins = np.histogram(avg_uhis, bins=50)
    max_counts, max_bins = np.histogram(max_uhis, bins=50)
    #tim_counts, tim_bins = np.histogram(max_time, bins=24)

    avg_counts = np.append(0, avg_counts)/ndays*365
    max_counts = np.append(0, max_counts)/ndays*365
    #tim_counts = np.append(0, tim_counts)/ndays*365

    ax[0].step(avg_bins, avg_counts, where='post', color=plt_colors[icity], label=city['name'])
    ax[1].step(max_bins, max_counts, where='post', color=plt_colors[icity], label=city['name'])
    #ax[2].step(tim_bins, tim_counts, where='post', color=plt_colors[icity], label=city['name'])

    ax[0].fill_between(avg_bins, avg_counts, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')
    ax[1].fill_between(max_bins, max_counts, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')
    #ax[2].fill_between(tim_bins, tim_counts, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')

    city['avg_counts'] = avg_counts
    city['avg_bins'  ] = avg_bins
    city['max_counts'] = max_counts
    city['max_bins'  ] = max_bins

ax[0].set_ylabel(r'nr. days / year')
ax[0].set_xlabel(r'<UHI>$_{\mathrm{day}}$')
ax[1].set_xlabel(r'max(UHI)$_{\mathrm{day}}$')
ax[0].set_ylim([0, 34])
ax[1].set_ylim([0, 49])
ax[1].legend(loc=0, fontsize='small')
ax[0].text(-.15, .95, '(a)', transform=ax[0].transAxes)
ax[1].text(-.15, .95, '(b)', transform=ax[1].transAxes)
plt.draw()
plt.savefig(imgsDir+'uhi_histograms.png')
plt.savefig(imgsDir+'uhi_histograms.pdf')

tval = 7
for city in cities:
    #N = np.sum(city['avg_counts'][np.where(city['avg_bins']>tval)[0]])
    N = np.sum(city['max_counts'][np.where(city['max_bins']>tval)[0]])
    N = N*n_years
    print('{0:10s}: {1:8.2f}'.format(city['name'], N))


#===============================================================================
# Figure
# Tropical nights
#
fig = plt.figure(4, figsize=(10,4.76)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(2, 1, num=fig.number, sharex=True, sharey=False)
scale=0.09
for icity, city in enumerate(cities):
    ax[0].bar(np.array(years) - scale*(len(cities)-1)/2 + scale*icity, city['sim_city_nights'],  width = scale, label=city['name'])
    ax[1].bar(np.array(years) - scale*(len(cities)-1)/2 + scale*icity, city['sim_rural_nights'], width = scale, label=city['name'])
    #ax[2].bar(np.array(years) - scale*(len(cities)-1)/2 + scale*icity, city['sim_city_nights']/city['sim_rural_nights'], width = scale, label=city['name'])
ax[0].set_ylabel('City centre')
ax[1].set_ylabel('Rural reference')
ax[1].set_xticks(years)
ax[0].set_yticks([0,25,50,75,100])
ax[1].set_yticks([0,25,50,75,100])
ax[0].grid(which='major', axis='y')
ax[1].grid(which='major', axis='y')
ax[0].set_axisbelow(True)
ax[1].set_axisbelow(True)
# Shrink the axes
for iax in range(len(ax)):
    box = ax[iax].get_position()
    ax[iax].set_position([box.x0, box.y0, box.width * 0.9, box.height])
# Put a legend to the right of the current axis
ax[0].legend(loc='center left', bbox_to_anchor=(1, 0), fontsize='small')
#ax[0].legend(fontsize='small')
plt.draw()
plt.savefig(imgsDir+'tropical_nights_atCities.pdf')
plt.savefig(imgsDir+'tropical_nights_atCities.png')

#===============================================================================
# Figure
# Time series of UHI w.r.t. sunset
#
from astral import LocationInfo
from astral.sun import sun
from datetime import timezone

for icity, city in enumerate(cities):

    I, J = city['cityIJ']
    a_city = LocationInfo(city['name'], region='Switzerland', timezone='Europe/Zurich', latitude=grid_lat[J,I], longitude=grid_lon[J,I])

    all_times  = np.array(city['times'])
    all_rural  = np.array(city['t2m_rural'])
    all_cities = np.array(city['t2m_city'])
    all_uhis   = all_cities - all_rural

    days = []
    delays = []

    maxTime = []
    maxIdx  = []

    # go day by day, find max uhi
    p0 = p1 = 0
    while p1 < len(all_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(all_times >= all_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(all_times)+1

        p2 = (p0+p1)//2

        #idx = np.argmax( all_uhis[p0:p1] )
        #maxUHI_time = all_times[p0:p1][idx]
        idx = np.argmax( all_uhis[p2:p1] )
        maxUHI_time = all_times[p2:p1][idx]
        sunset_time = sun(a_city.observer, date=all_times[p0])['sunset']

        days.append( all_times[p0] )
        delays.append( (maxUHI_time.astimezone(sunset_time.tzinfo) - sunset_time).total_seconds()/3600 )

        maxTime.append( maxUHI_time )
        maxIdx.append( p2+idx )

    days   = np.array(days)
    delays = np.array(delays)
    maxTime = np.array(maxTime)
    maxIdx  = np.array(maxIdx)

fig = plt.figure(5, figsize=(12,8)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(2, 1, num=fig.number, sharex=True, sharey=False)
ax[0].plot(days,delays)
ax[1].plot(all_times, all_uhis)
ax[1].plot(maxTime, all_uhis[maxIdx], 'o')
plt.draw()

#===============================================================================
# Figure
# Time series of UHI SEASON DIEL COMPOSITE
#

# compute average sunset times
from astral import LocationInfo
from astral.sun import sun
from datetime import timezone
for icity, city in enumerate(cities):

    I, J = city['cityIJ']
    a_city = LocationInfo(city['name'], region='Switzerland', timezone='Europe/Zurich', latitude=grid_lat[J,I], longitude=grid_lon[J,I])
    date = datetime(2017,1,1)
    city['sunset']  = np.zeros((4,2))
    city['sunrise'] = np.zeros((4,2))
    while date < datetime(2018,1,1):
        sunset_time  = sun(a_city.observer, date=date)['sunset']
        sunrise_time = sun(a_city.observer, date=date)['sunrise']
        month = date.month
        if   (month >=12) | (month<=2):
            season = 0
        elif (month >  2) & (month<=5):
            season = 1
        elif (month >  5) & (month<=8):
            season = 2
        else:
            season = 3
        city['sunset'][season, 0] += sunset_time.hour + sunset_time.minute/60
        city['sunset'][season, 1] += 1
        city['sunrise'][season, 0] += sunrise_time.hour + sunrise_time.minute/60
        city['sunrise'][season, 1] += 1
        date += relativedelta(days=1)
    city['sunset'] [:, 0] /= city['sunset'] [:, 1]
    city['sunrise'][:, 0] /= city['sunrise'][:, 1]


fig = plt.figure(6, figsize=(9.5,13.5)); plt.clf(); plt.show(block=False)
sharey=True
fig, ax = plt.subplots(len(cities), 4, num=fig.number, sharex=True, sharey=sharey)
for icity, city in enumerate(cities):

    times  = np.arange(24)
    uhis   = np.zeros((4, 24))
    counts = np.zeros((4, 24))
    v_uhis = np.zeros((4, 24))

    a_times  = np.array(city['times'])
    a_rural  = np.array(city['t2m_rural'])
    a_cities = np.array(city['t2m_city'])
    a_uhis   = a_cities - a_rural

    # Mean
    for idx in range(len(a_times)):
        month = a_times[idx].month
        if   (month >=12) | (month<=2):
            season = 0
        elif (month >  2) & (month<=5):
            season = 1
        elif (month >  5) & (month<=8):
            season = 2
        else:
            season = 3
        hour = a_times[idx].hour
        uhis  [season, hour] += a_uhis[idx]
        counts[season, hour] += 1
    uhis = uhis / counts

    # Variance
    for idx in range(len(a_times)):
        month = a_times[idx].month
        if   (month >=12) | (month<=2):
            season = 0
        elif (month >  2) & (month<=5):
            season = 1
        elif (month >  5) & (month<=8):
            season = 2
        else:
            season = 3
        hour = a_times[idx].hour
        v_uhis[season, hour] += (a_uhis[idx] - uhis[season,hour])**2
    v_uhis = (v_uhis / (counts-1))**0.5

    for season in range(4):
        ax[icity,season].plot(times, uhis[season,:], label='UHI')
        ax[icity,season].fill_between(times, uhis[season,:]-v_uhis[season,:], uhis[season,:]+v_uhis[season,:],alpha=0.5)
        #
        if (sharey==True) or (sharey=='row'):
            ax[icity,season].plot([0,24], [0,0], '-k', alpha=0.3)
        else:
            ylims = ax[icity,season].get_ylim()
            ax[icity,season].plot([0,24], [0,0], '-k', alpha=0.3)
            ax[icity,season].set_ylim(ylims) # reset ylims
        #
        uhi_delta = uhis[season,:].max() - uhis[season,:].min()
        if sharey==True:
            ax[icity,season].text(1, -0.8,                 r'$\Delta = $ {0:.1f}'.format(uhi_delta), fontsize='small', ha='left', va='center')
        else:
            ax[icity,season].text(1, uhis[season,:].min(), r'$\Delta = $ {0:.1f}'.format(uhi_delta), fontsize='small', ha='left', va='center')
        ax[icity,season].set_xlim([0,23])
        ax[icity,season].set_xticks([0, 6, 12, 18])
        #
        # add sunset
        ylims = ax[icity,season].get_ylim()
        if sharey==True:
            ylims = [-1.2,+2.8]
        ax[icity,season].plot(2*[city['sunrise'][season,0]], [ylims[0]+.2, ylims[1]-.2], '--k', alpha=0.3)
        ax[icity,season].plot(2*[city['sunset'][season,0]],  [ylims[0]+.2, ylims[1]-.2], '--k', alpha=0.3)

    ax[icity,0].set_ylabel(city['name'])

    # save
    city[ 'times_uhi'] =  times
    city[  'uhis_uhi'] =   uhis
    city['v_uhis_uhi'] = v_uhis

ax[-1,0].set_xlabel('Winter [DJF]')
ax[-1,1].set_xlabel('Spring [MAM]')
ax[-1,2].set_xlabel('Summer [JJA]')
ax[-1,3].set_xlabel('Autumn [SON]')
#ax[-1,-1].legend(fontsize='small')
plt.draw()
if sharey==True:
    plt.savefig(imgsDir+'cities_UHI_seasonComposite_b.png')
    plt.savefig(imgsDir+'cities_UHI_seasonComposite_b.pdf')
elif sharey=='row':
    plt.savefig(imgsDir+'cities_UHI_seasonComposite_c.png')
    plt.savefig(imgsDir+'cities_UHI_seasonComposite_c.pdf')
else:
    plt.savefig(imgsDir+'cities_UHI_seasonComposite.png')
    plt.savefig(imgsDir+'cities_UHI_seasonComposite.pdf')

#===============================================================================
# Figure
# Time series of Temperature to understand the figure above SEASON DIEL COMPOSITE
#
fig = plt.figure(7, figsize=(9.5,13.5)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(len(cities), 4, num=fig.number, sharex=True, sharey=False)
for icity, city in enumerate(cities):

    times    = np.arange(24)
    t_cities = np.zeros((4, 24))
    t_rurals = np.zeros((4, 24))
    counts   = np.zeros((4, 24))
    v_cities = np.zeros((4, 24))
    v_rurals = np.zeros((4, 24))

    a_times  = np.array(city['times'])
    a_rurals = np.array(city['t2m_rural'])
    a_cities = np.array(city['t2m_city'])

    # Mean
    for idx in range(len(a_times)):
        month = a_times[idx].month
        if   (month >=12) | (month<=2):
            season = 0
        elif (month >  2) & (month<=5):
            season = 1
        elif (month >  5) & (month<=8):
            season = 2
        else:
            season = 3
        hour = a_times[idx].hour
        t_cities[season, hour] += a_cities[idx]
        t_rurals[season, hour] += a_rurals[idx]
        counts  [season, hour] += 1
    t_cities = t_cities / counts
    t_rurals = t_rurals / counts

    # Variance
    for idx in range(len(a_times)):
        month = a_times[idx].month
        if   (month >=12) | (month<=2):
            season = 0
        elif (month >  2) & (month<=5):
            season = 1
        elif (month >  5) & (month<=8):
            season = 2
        else:
            season = 3
        hour = a_times[idx].hour
        v_cities[season, hour] += (a_cities[idx] - t_cities[season,hour])**2
        v_rurals[season, hour] += (a_rurals[idx] - t_rurals[season,hour])**2
    v_cities = (v_cities / (counts-1))**0.5
    v_rurals = (v_rurals / (counts-1))**0.5

    for season in range(4):
        ax[icity,season].plot(times, t_cities[season,:], label='UHI')
        ax[icity,season].fill_between(times, t_cities[season,:]-v_cities[season,:], t_cities[season,:]+v_cities[season,:],alpha=0.5)
        ax[icity,season].plot(times, t_rurals[season,:], label='UHI')
        ax[icity,season].fill_between(times, t_rurals[season,:]-v_rurals[season,:], t_rurals[season,:]+v_rurals[season,:],alpha=0.5)
        ax[icity,season].set_xlim([0,23])
        ax[icity,season].set_xticks([0, 6, 12, 18])

    ax[icity,0].set_ylabel(city['name'])

ax[-1,0].set_xlabel('Winter [DJF]')
ax[-1,1].set_xlabel('Spring [MAM]')
ax[-1,2].set_xlabel('Summer [JJA]')
ax[-1,3].set_xlabel('Autumn [SON]')
plt.draw()

plt.figure(101)
plt.clf()
plt.plot(times, t_cities[season,:], '-+', label='UHI')
plt.fill_between(times, t_cities[season,:]-v_cities[season,:], t_cities[season,:]+v_cities[season,:],alpha=0.5)
plt.plot(times, t_rurals[season,:], '-+', label='UHI')
plt.fill_between(times, t_rurals[season,:]-v_rurals[season,:], t_rurals[season,:]+v_rurals[season,:],alpha=0.5)
plt.xlim([0,23])
plt.xticks([0, 6, 12, 18])
plt.draw()

plt.figure(102); plt.clf(); plt.show(block=False)
a_times  = np.array(city['times'])
a_rurals = np.array(city['t2m_rural'])
a_cities = np.array(city['t2m_city'])
plt.plot(a_times, a_cities, '-', label='UHI')
plt.plot(a_times, a_rurals, '-', label='UHI')
plt.draw()

#===============================================================================
# Compare with
# Kubilay et al., ‘Advancement in Urban Climate Modelling at Local Scale’.
#
city = cities[0]
a_times   = np.array(city['times'])
a_rural   = np.array(city['t2m_rural'])
a_cities  = np.array(city['t2m_city'])
a_uhis    = a_cities - a_rural

# UHI 2019 June 24 to July 1
#
idxs0 = set(np.argwhere(a_times >= datetime(2019, 6, 24)).T[0])
idxs1 = set(np.argwhere(a_times <= datetime(2019, 7,  1)).T[0])
idxs  = list(set.intersection(idxs0, idxs1))
c_uhi = a_uhis[idxs]

#===============================================================================
# Compare with
# Gubler et al., ‘Evaluation and Application of a Low-Cost Measurement Network to Study Intra-Urban Temperature Differences during Summer 2018 in Bern, Switzerland’.
#
city = cities[4]
a_times   = np.array(city['times'])
a_rural   = np.array(city['t2m_rural'])
a_cities  = np.array(city['t2m_city'])
a_uhis    = a_cities - a_rural

# UHI 2018 July 30 to August 7
#
idxs0 = set(np.argwhere(a_times >= datetime(2018, 7, 30)).T[0])
idxs1 = set(np.argwhere(a_times <= datetime(2018, 8,  7)).T[0])
idxs  = list(set.intersection(idxs0, idxs1))
c_uhi = a_uhis[idxs]


#===============================================================================
# Figure
# Dry, clear, calm days
#
tol_prec = 0.2 # max total daily rainfall [mm]
tol_clct = 0.2 # cloud coverage / percent of passing radiation
tol_winM = 3.0 # max night magnitude [m/s]

fig = plt.figure(8, figsize=(20,3.62)); plt.clf(); plt.show(block=False)
fig, ax = plt.subplots(1, 5, num=fig.number, sharex=True, sharey=True)
for icity, city in enumerate(cities):

    all_times  = np.array(city['times'])
    all_rural  = np.array(city['t2m_rural'])
    all_cities = np.array(city['t2m_city'])
    all_prec   = np.array(city['avg_prec'])
    all_winM   = np.array(city['avg_winM'])
    all_clct   = np.array(city['avg_clct'])
    #
    all_uhis   = all_cities - all_rural

    days0 = []
    days1 = []
    days2 = []
    days3 = []
    uhis0 = []
    uhis1 = []
    uhis2 = []
    uhis3 = []

    p0 = p1 = 0
    while p1 < len(all_times)-24:
        p0 = p1
        try:
            p1 = np.argwhere(all_times >= all_times[p0]+relativedelta(days=1))[0][0]
        except:
            p1 = len(all_times)+1

        today = all_times[p0]
        uhi   = np.mean( all_uhis[p0:p1] )

        # all days
        days0.append( today )
        uhis0.append( uhi   )
        # conditionals
        if np.sum( all_prec[p0:p1] ) <= tol_prec:
            # dry
            days1.append( today )
            uhis1.append( uhi   )
            if np.mean( all_clct[p0:p1] ) <= tol_clct:
                # dry,clear
                days2.append( today )
                uhis2.append( uhi   )
                if np.max( all_winM[p0:p1] ) <= tol_winM:
                    # dry,clear,calm
                    days3.append( today )
                    uhis3.append( uhi   )

    days0 = np.array(days0)
    days1 = np.array(days1)
    days2 = np.array(days2)
    days3 = np.array(days3)
    uhis0 = np.array(uhis0)
    uhis1 = np.array(uhis1)
    uhis2 = np.array(uhis2)
    uhis3 = np.array(uhis3)

    #bins = np.linspace(-0.5, 5.5, 30)
    #bins = 30
    bins = np.arange(-10,10,0.125)
    counts0, bins0 = np.histogram(uhis0, bins=bins)
    counts1, bins1 = np.histogram(uhis1, bins=bins)
    counts2, bins2 = np.histogram(uhis2, bins=bins)
    counts3, bins3 = np.histogram(uhis3, bins=bins)

    print('{:10s}: {:5d} {:5d} {:5d} {:5d}'.format(city['name'], np.sum(counts0), np.sum(counts1), np.sum(counts2), np.sum(counts3)))

    # normalisation
    counts0 = np.append(0, counts0)/np.sum(counts0)
    counts1 = np.append(0, counts1)/np.sum(counts1)
    counts2 = np.append(0, counts2)/np.sum(counts2)
    counts3 = np.append(0, counts3)/np.sum(counts3)


    ax[0].step(bins0, counts0, where='post', color=plt_colors[icity], label=city['name'])
    ax[1].step(bins1, counts1, where='post', color=plt_colors[icity], label=city['name'])
    ax[2].step(bins2, counts2, where='post', color=plt_colors[icity], label=city['name'])
    ax[3].step(bins3, counts3, where='post', color=plt_colors[icity], label=city['name'])

    ax[0].fill_between(bins0, counts0, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')
    ax[1].fill_between(bins1, counts1, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')
    ax[2].fill_between(bins2, counts2, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')
    ax[3].fill_between(bins3, counts3, step='post', color=plt_colors[icity], alpha=0.2, label='_nolegend_')

    ax[0].text(0.03+(icity%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins0[np.argmax(counts0)]), color=plt_colors[icity], transform=ax[0].transAxes)
    ax[1].text(0.03+(icity%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins1[np.argmax(counts1)]), color=plt_colors[icity], transform=ax[1].transAxes)
    ax[2].text(0.03+(icity%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins2[np.argmax(counts2)]), color=plt_colors[icity], transform=ax[2].transAxes)
    ax[3].text(0.03+(icity%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins3[np.argmax(counts3)]), color=plt_colors[icity], transform=ax[3].transAxes)

    city['counts0'] = counts0; city['bins0'] = bins0
    city['counts1'] = counts1; city['bins1'] = bins1
    city['counts2'] = counts2; city['bins2'] = bins2
    city['counts3'] = counts3; city['bins3'] = bins3

# Extra city
if True:
    weather_colors = ['#A7A7A7', '#F3C960', '#BCEC98', '#CAEAF9']
    weather_colors = ['#6D6D6D', '#BF9A42', '#8EA044', '#7CADC9']
    icity=0
    city=cities[icity]
    counts0=city['counts0']; bins0=city['bins0']
    counts1=city['counts1']; bins1=city['bins1']
    counts2=city['counts2']; bins2=city['bins2']
    counts3=city['counts3']; bins3=city['bins3']
    ax[4].step(bins0, counts0, where='post', color=weather_colors[0], label='all')
    ax[4].step(bins1, counts1, where='post', color=weather_colors[1], label='dry')
    ax[4].step(bins2, counts2, where='post', color=weather_colors[2], label='dry,clear')
    ax[4].step(bins3, counts3, where='post', color=weather_colors[3], label='dry,clear,calm')
    ax[4].fill_between(bins0, counts0, step='post', color=weather_colors[0], alpha=0.4, label='_nolegend_')
    ax[4].fill_between(bins1, counts1, step='post', color=weather_colors[1], alpha=0.4, label='_nolegend_')
    ax[4].fill_between(bins2, counts2, step='post', color=weather_colors[2], alpha=0.4, label='_nolegend_')
    ax[4].fill_between(bins3, counts3, step='post', color=weather_colors[3], alpha=0.4, label='_nolegend_')
    ax[4].text(0.03+(0%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins0[np.argmax(counts0)]), color=weather_colors[0], transform=ax[4].transAxes)
    ax[4].text(0.03+(1%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins1[np.argmax(counts1)]), color=weather_colors[1], transform=ax[4].transAxes)
    ax[4].text(0.03+(2%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins2[np.argmax(counts2)]), color=weather_colors[2], transform=ax[4].transAxes)
    ax[4].text(0.03+(3%5)*0.18, 0.88-int(icity>=5)*0.09, '{:4.2f}'.format(bins3[np.argmax(counts3)]), color=weather_colors[3], transform=ax[4].transAxes)
    ax[4].legend(loc='center left', bbox_to_anchor=(0, 0.75), fontsize='small')


ax[0].text(0.91, 0.93, '(a)', transform=ax[0].transAxes)
ax[1].text(0.91, 0.93, '(b)', transform=ax[1].transAxes)
ax[2].text(0.91, 0.93, '(c)', transform=ax[2].transAxes)
ax[3].text(0.91, 0.93, '(d)', transform=ax[3].transAxes)
ax[4].text(0.91, 0.93, '(e)', transform=ax[4].transAxes)

ax[0].set_ylabel(r'Frequency distribution')
ax[0].set_xlabel(r'all')
ax[1].set_xlabel(r'dry')
ax[2].set_xlabel(r'dry,clear')
ax[3].set_xlabel(r'dry,clear,calm')
ax[4].set_xlabel(r'{:s}'.format(city['name']))
ax[0].set_xlim([-1.5, 3.5])
ax[0].set_ylim([0, .35])
#ax[-1].legend(loc=0, fontsize='small')
#ax[0].text(-.15, .95, '(c)', transform=ax[0].transAxes)
#ax[1].text(-.15, .95, '(d)', transform=ax[1].transAxes)
plt.draw()
# # Shrink the axes
# for iax in range(len(ax)):
#     box = ax[iax].get_position()
#     ax[iax].set_position([box.x0, box.y0, box.width * 0.9, box.height])
# Put a legend to the right of the current axis
ax[3].legend(loc='center left', bbox_to_anchor=(2.2, 0.5))#, fontsize='small')
plt.draw()

plt.savefig(imgsDir+'uhi_histograms_c.png')
plt.savefig(imgsDir+'uhi_histograms_c.pdf')