#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct  1 16:36:42 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 as Pi, cos, sin
import cftime
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
year = 2000
#year_new = 2000# use new time axis in year 2000 
negative_natural = True # rotate local (x,y) coordinate to negative natural coordinate (TC moves to the x-negative direction)
"""
data_name = 'precip'
scale_factor = 24*3600
units = 'mm/day'

data_name = 'msl'
scale_factor = 1e-2 #Pa -> hPa
units = 'hPa'
"""
data_name = 'v10'
scale_factor = 1
units = 'm/s'

odata_name = f'{data_name}_xy'
ofile = f'{odata_name}_{year:04d}.nc'

R = 6370 #km, earth radius
L1 = 2*Pi*R/360 #km, distance of one degree longitude at the equator (or latitude at any location)
x = np.arange(-500, 501, 10)
y = np.arange(-500, 501, 10)
x = xr.DataArray(x, dims='x', coords=[x], attrs=dict(units='km', long_name='local x'))
y = xr.DataArray(y, dims='y', coords=[y], attrs=dict(units='km', long_name='local y'))
if negative_natural:
    x = x.assign_attrs(long_name='negative natural s')
    y = y.assign_attrs(long_name='negative natural n')

print('loading data ...')
# composite variable
ifiles = f'/tigress/wenchang/data/era5/raw/10m_v_component_of_wind/era5.10m_v_component_of_wind.{year:04d}-??.nc'
da = xr.open_mfdataset(ifiles)[data_name] \
    .pipe(lambda x: x*scale_factor) \
    .assign_attrs(units=units) \
    .rename(longitude='lon', latitude='lat')
da.load()
# tc tracks
ds = xr.open_dataset('/tigress/wenchang/data/ibtracs/v04r00/analysis/v2/IBTrACS.ALL.v04r00.tracksByYear.1980-2019.nc') \
    .sel(year=year)
ds.load()

time_adjust = False # needed for GFDL model output and TC tracks since the hour information is saved unconventionally (hour 24 instead of 00 is used)
if time_adjust:
    # time axis adjustment
    # change year so that time range is 2000-01-01 00:00:00 to 2001-01-01 00:00:00 
    time = da.time
    time_new = [cftime.DatetimeGregorian(yr-year+2000, month, day, hour)
        for yr,month,day,hour in zip(time.dt.year.values, time.dt.month.values, time.dt.day.values, time.dt.hour.values)]
    da = da.assign_coords(time=time_new)
    # adjust time stamp if no track points on Dec 31: shift da time index one day back from Mar 1st to the end of year to match the month/day/hour from tracks data (the mismatch may arise during TC tracking where non-leap year is treated as leap year
    if not ds.hour.where((ds.month==12)&(ds.day==31)).max()>0: 
        time = da.time
        times = np.array([f'{yr:04d}-{month:02d}-{day:02d} {hour:02d}:00:00'
            for yr,month,day,hour in zip(time.dt.year.values, time.dt.month.values, time.dt.day.values, time.dt.hour.values)])
        time_new = da.indexes['time'].where( times<'2000-03-01 00:00:00', other=da.indexes['time'].shift(-1, 'D') )
        da = da.assign_coords(time=time_new)
    # adjust hour to avoid the problem of hour 24:00:00 in tracks dasta (only 00:00:00 from da)
    da = da.assign_coords(time=da.indexes['time'].shift(-1, 'H')) # so that no 24:00:00 (not recognized by python)

# interpolation loop
zz = np.zeros((ds.storm.size, ds.stage.size, y.size, x.size)) + np.nan
angle_tcmotion = np.zeros((ds.storm.size, ds.stage.size)) + np.nan
vtcmotion = np.zeros((ds.storm.size, ds.stage.size)) + np.nan
for istorm in range(ds.storm.size):
    #print('istorm:', istorm)
    # stop if no more tracks data
    if np.isnan(ds.isel(storm=istorm, stage=0).lon.item()):
        break
    for istage in range(ds.stage.size):
        #s = f'year: {year:04d}; storm #: {istorm+1:03d}; step #: {istage:03d}'
        #print(s)
        # current track point
        ds0 = ds.isel(storm=istorm, stage=istage)
        lon0 = ds0.lon.item()
        # stop if not a valid track point (nan value)
        if np.isnan(lon0):
            s = f'year: {year:04d}; storm #: {istorm+1:03d}; step #: {istage:03d}'
            print(s)
            break
        lat0 = ds0.lat.item()
        month0 = int( ds0.month.item() )
        day0 = int( ds0.day.item() )
        hour0 = int( ds0.hour.item() )
        if time_adjust:
            hour0 = hour0 - 1 # minus 1 to be consistent with da.time
        
        # track point time stamp
        if time_adjust:
            t0 = f'2000-{month0:02d}-{day0:02d}T{hour0:02d}'
        else:
            t0 = f'{year:04d}-{month0:02d}-{day0:02d}T{hour0:02d}'

        # local coordinates
        # rotate local coordinate so that TC moves to the left (direction of negative x axis): negative natural coordinate
        if negative_natural:
            # get dlat and dlon
            if istage == 0: # first track point: forward difference
                ds_next = ds.isel(storm=istorm, stage=1)
                dlat = ds_next.lat.item() - ds0.lat.item()
                dlon = ds_next.lon.item() - ds0.lon.item()
            elif np.isnan(ds.isel(storm=istorm, stage=istage+1).lon.item()): # last track point: backward difference
                ds_prev = ds.isel(storm=istorm, stage=istage-1)
                dlat = ds0.lat.item() - ds_prev.lat.item()
                dlon = ds0.lon.item() - ds_prev.lon.item()
            else: #central difference
                ds_next = ds.isel(storm=istorm, stage=istage+1)
                ds_prev = ds.isel(storm=istorm, stage=istage-1)
                dlat = ( ds_next.lat.item() - ds_prev.lat.item() )/2
                dlon = ( ds_next.lon.item() - ds_prev.lon.item() )/2
            angle_tcmotion[istorm, istage] = np.angle(dlon*cos(lat0*Pi/180) + dlat*1j, deg=True)# angle (in degree) between TC motion and zonal direction
            dx, dy, dt = L1*dlon*cos(lat0*Pi/180)*1000, L1*dlat*1000, 6*3600 # in units of m, m, s
            vtcmotion[istorm, istage] = (dx**2 + dy**2)**0.5/dt # tc motion speed: m/s
            theta = np.angle(-dlon*cos(lat0*Pi/180) - dlat*1j) # angle between x-negative in rotating axis and zonal direction
            # x-negative natural coordinate (x,y) projected onto local (x_, y_) coordinate
            x_ = x*cos(theta) - y*sin(theta)
            y_ = x*sin(theta) + y*cos(theta)
            lat_xy = (y_+x_)*0 + lat0 + y_/L1 
            lon_xy = (y_+x_)*0 + lon0 + x_/(L1*cos(lat_xy*Pi/180))
        else:
            lat_xy = (y+x)*0 + lat0 + y/L1 
            lon_xy = (y+x)*0 + lon0 + x/(L1*cos(lat_xy*Pi/180))

        # interpolation from regular lon/lat grids onto local negative natural coordinates
        da0 = da.sel(time=t0).squeeze() \
            .interp(lon=lon_xy, lat=lat_xy) \
            .drop(['lon', 'lat', 'time'])
        da0 = da0.transpose('y', 'x')

        zz[istorm, istage, :, :] = da0.values

# ndarray to DataArray
da_ = xr.DataArray(zz, dims=['storm', 'stage', 'y', 'x'], 
    coords=[ds.storm, ds.stage, y, x], attrs=dict(units=units))

# save to nc
print('saving to', ofile, '...')
ds[odata_name] = da_
if negative_natural:
    angle_tcmotion = xr.DataArray(angle_tcmotion, dims=['storm', 'stage'], 
        coords=[ds.storm, ds.stage], 
        attrs=dict(units='deg', long_name='angle between TC motion and zonal east'))
    ds['angle_tcmotion'] = angle_tcmotion
    vtcmotion = xr.DataArray(vtcmotion, dims=['storm', 'stage'],
        coords=[ds.storm, ds.stage],
        attrs=dict(units='m/s', long_name='TC motion speed'))
    ds['vtcmotion'] = vtcmotion
ds.to_netcdf(ofile,
    encoding={odata_name: {'zlib': True, 'complevel': 1, 'dtype': 'float32'}})
print('[saved]:', ofile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    da_.mean(['storm', 'stage']).plot()
    #plt.savefig(figname)
    #print('[saved]:', figname)
    
    tt.check(f'**Done**')
    plt.show()