#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Nov 17 14:56:34 EST 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
#import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
latlons = {
    'Buenos Aires': (-34.6037, 360 - 58.3816),
    'Delhi': (28.7041, 77.1025),
    'London': (51.5074, 360 - 0.1278),
    'Los Angeles': (34.0522, 360 - 118.2437),
    'Melbourne': (-37.8136, 144.9631),
    'Miami': (25.7617, 360 - 80.1918),
    'New York': (40.7128, 360 - 74.0060),
    'Phoenix': (33.4484, 360 - 112.0740),
    'Sinaloa': (25.1721, 360 - 107.4795),
    'Wuhan': (30.5928, 114.3055),
    }
location = ('Buenos Aires', 'Delhi', 'London', 'Los Angeles', 'Melbourne', 'Miami', 'Phoenix', 'New York', 'Sinaloa', 'Wuhan')[-2]
lat, lon = latlons[location]
#lat, lon = 33.4484, 360 - 112.0740#40.7128, 360 - 74.0060 # New York lat/lon
q2m = xr.open_dataarray('/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/clim/era5.q2m.w1clim1981-2010.nc')
t2m = xr.open_dataarray('/tigress/wenchang/data/era5/analysis_wy/2m_temperature/daily/clim/era5.t2m.w1clim1981-2010.nc')
long_name = t2m.attrs['long_name']
t2m = t2m.pipe(lambda x: x - 273.15).assign_attrs(units='$^\circ$C', long_name=long_name)
uv = xr.open_dataarray('/tigress/wenchang/data/era5/analysis_wy/mean_surface_downward_uv_radiation_flux/daily/clim/era5.msdwuvrf.w1clim1981-2010.nc')

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figname = __file__.replace('.py', f'_{location.replace(" ", "-")}_{tt.today()}.png')
    figname = os.path.join(__file__.replace('.py', ''), figname)
    fig, axes = plt.subplots(2, 1, figsize=(7, 6))

    ax0 = axes[0]
    q2m.sel(longitude=lon, latitude=lat, method='nearest').plot(color='C0', marker='o', label='q2m', ax=ax0)
    ax0_twinx = ax0.twinx()
    t2m.sel(longitude=lon, latitude=lat, method='nearest').plot(color='C3', marker='^',  label='t2m', ax=ax0_twinx)

    if lat > 0:
        ax0.legend(loc='upper left')
        ax0_twinx.legend(loc='upper right')
    else:
        ax0.legend(loc='lower left')
        ax0_twinx.legend(loc='lower right')
    ax0.set_title(location, loc='left')
    ax0_twinx.set_title('')

    ax1 = axes[1]
    q2m.sel(longitude=lon, latitude=lat, method='nearest').plot(color='C0', marker='o', label='q2m', ax=ax1)
    ax1_twinx = ax1.twinx()
    uv.sel(longitude=lon, latitude=lat, method='nearest').plot(color='C1', marker='*',  label='uv', ax=ax1_twinx)

    if lat > 0:
        ax1.legend(loc='upper left')
        ax1_twinx.legend(loc='upper right')
    else:
        ax1.legend(loc='lower left')
        ax1_twinx.legend(loc='lower right')
    ax1.set_title(location, loc='left')
    ax1_twinx.set_title('')

    fig.tight_layout()
    plt.savefig(figname)
    print('[saved]:', figname)

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