#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Oct 14 14:37:12 EDT 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
from numpy import pi, cos, sin
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
load_data = False
if load_data:
    print('loading data ...')
    # load data
    ds = xr.open_dataset('test_windspeed_nn.nc')
    da = ds.windspeed_xy

# target polar coordinate
r = np.arange(0, 501, 10)
theta = np.arange(0, 361, 5)
# wrap into DataArray
r = xr.DataArray(r, dims='r', attrs=dict(units='km', long_name='radius'))
r = r.assign_coords(r=r)
theta = xr.DataArray(theta, dims='theta', attrs=dict(units='deg', long_name='angle'))
theta = theta.assign_coords(theta=theta)

#convert to Cartesian coordinate and interpolate
x_rtheta = r*cos(theta/180*pi)
y_rtheta = r*sin(theta/180*pi)
# interpolate
da_ = da.interp(x=x_rtheta, y=y_rtheta)
da_ = da_.rename(da_.name.replace('xy', 'rtheta'))

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    close('all')
    fig = plt.figure()
    ax = fig.add_subplot(projection='polar')
    im = ax.contourf(theta/180*pi, r, x_rtheta, cmap='turbo')
    plt.colorbar(im, label='x')

    fig = plt.figure()
    ax = fig.add_subplot(projection='polar')
    im = ax.contourf(theta/180*pi, r, y_rtheta, cmap='turbo')
    plt.colorbar(im, label='y')

    figure()
    da_.where(ds.tm>0).mean(['stage', 'storm', 'theta']).plot(label='all WC track points', color='k')
    da_.where(ds.tm>0).where(ds.lat>0).mean(['stage', 'storm']).sel(theta=slice(0, 180)).mean('theta').plot(label='NH right of TC motion')
    da_.where(ds.tm>0).where(ds.lat>0).mean(['stage', 'storm']).sel(theta=slice(180, 360)).mean('theta').plot(label='NH left of TC motion')
    da_.where(ds.tm>0).where(ds.lat<0).mean(['stage', 'storm']).sel(theta=slice(0, 180)).mean('theta').plot(label='SH right of TC motion', ls='--')
    da_.where(ds.tm>0).where(ds.lat<0).mean(['stage', 'storm']).sel(theta=slice(180, 360)).mean('theta').plot(label='SH left of TC motion', ls='--')
    plt.ylabel('wind speed [m/s]')
    plt.legend()
    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    plt.savefig(figname)
    print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()