# -*- coding: utf-8 -*-
"""
Created on Fri Jan 29 10:09:27 2021

@author: ppantina
"""
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.dates as mdate

def L1B_sub_along_cross_winds (radName, flightDate, L1B, L1B_uwind, L1B_vwind, L1B_vel, cmap_n0q, deg_offset):
    ##Isolate vars
    uwnd  = L1B_uwind
    vwnd  = L1B_vwind
    time  = L1B['Time/Data/TimeUTC']          [:]
    height= L1B['Navigation/Data/Height']     [:]
    head  = L1B['Navigation/Data/Heading']    [:]
    roll  = L1B['Navigation/Data/Roll']       [:]
    pitch = L1B['Navigation/Data/Pitch']      [:]
    hrange= L1B['Products/Information/Range'] [:]
    lon   = L1B['Navigation/Data/Longitude']  [:]    
    lat   = L1B['Navigation/Data/Latitude']   [:]    
    altitude = np.zeros((len(height), len(hrange)))
    for i, j in enumerate(altitude): #for all times.
        
        ##The height is just aircraft height - range
        altitude[i]= height[i] - hrange
    
    ##Convert uwind and vwind to a wind direction, with 0 as north.
    ##Need "90" here to translate away from 0 at y-axis.
    wdir = (90-np.rad2deg(np.arctan2(vwnd,uwnd)))%360
    
    ##Calculate wind speed from vectors
    wspd = np.sqrt(uwnd**2 + vwnd**2)
    
    ##Calculate a heading-relative wind direction. 
    ##Heading is from [-180,180] but theta will be [0,360]
    theta =  ((wdir.T - head)%360).T
    
    ##Define some storage arrays
    t_len = uwnd.shape[0]
    z_len = uwnd.shape[1]
    cross = np.zeros_like(uwnd, dtype = np.float32)*np.nan
    along = np.zeros_like(uwnd, dtype = np.float32)*np.nan
    
    ##Main calculation
    for tt in range(t_len): #for all time
        if tt %10000 == 0: print ('...Step ', tt, ' of ', t_len-1)
        for zz in range(z_len): #for all z
            ##Convert theta to radians
            theta_rad = np.deg2rad(theta[tt,zz])   
            
            ##Calculate cross/along track winds using basic geometry.
            ##CROSS is in x-direction, ALONG is in y-direction.
            cross[tt,zz] = wspd[tt,zz] * np.sin(theta_rad)
            along[tt,zz] = wspd[tt,zz] * np.cos(theta_rad)
    
  
    sin_roll  =   np.sin(np.deg2rad(roll + deg_offset))
    sin_pitch =  -np.sin(np.deg2rad(pitch))
    vel_corr_cross  = (cross.T * sin_roll).T
    vel_corr_along  = (along.T * sin_pitch).T
        
    
    '''
    secs = mdate.epoch2num(time)
    X, Y = np.meshgrid(secs, altitude[10000,:])
    
    date_fmt = '%H:%M'
    
    # Use a DateFormatter to set the data to the correct format.
    date_formatter = mdate.DateFormatter(date_fmt)
    
    fig = plt.figure(figsize = (14,8.5))
    ax= fig.add_subplot(311)
    plot1 = ax.imshow(L1B_vel.T + 0*vel_corr_cross.T, vmin= -1.5, vmax = 1.5, cmap = 'jet', aspect = 'auto', interpolation = 'nearest')
    cmap_n0q.set_under('w')
    ax.grid(True)
    ax.set_ylabel('Range (m)')
    ax.set_title('Velocity')
    ax.set_title(flightDate)
    ax= fig.add_subplot(312)
    plot2 = ax.imshow(L1B_vel.T -1*vel_corr_cross.T -1*vel_corr_along.T, vmin= -1.5, vmax = 1.5, cmap = 'jet', aspect = 'auto', interpolation = 'nearest')
    cmap_n0q.set_under('w')
    plt.colorbar(plot2, label = 'm/s', orientation = 'horizontal', shrink = 0.15)
    ax.grid(True)
    ax.set_ylabel('Range (m)')
    
    ax = fig.add_subplot(313)
    ax.plot(secs, (L1B_vel.T + 0*vel_corr_cross.T)[300,:], 'r.', markersize = 1, label = 'no corr')
    ax.plot(secs, (L1B_vel.T + 1*vel_corr_cross.T)[300,:], 'b.', markersize = 1, label = 'corr')
    plt.legend(prop={'size': 6})
    ax.set_ylabel('Mid-level Doppler Velocities')
    ax.set_xlabel('Time')
    ax.set_ylim(-2,2)
    ax.grid(True)
    plt.savefig('./plots/' + radName + '_' + flightDate + '_dopplercorrection_hrrr.png')
#    plt.show()
    
    fig = plt.figure(figsize = (11,8.5))
    plt.subplots_adjust(wspace=.1, hspace=.1)
    ax= fig.add_subplot(221)
    ax.plot(secs, wdir[:,400],         'C0.', markersize = 1, label = 'WIND DIR (deg)' )
    ax.plot(secs,cross[:,400], color = 'C1', label = 'CROSS WIND (m/s)')
    ax.plot(secs,along[:,400], color = 'C2', label = 'ALONG WIND (m/s)')
    ax.set_title(flightDate+'\nMetrics')
    ax.get_xaxis().set_visible(False)
    ax.set_ylabel('Absolute')
    ax.grid(True)
    plt.legend(prop={'size': 6})
    ax= fig.add_subplot(222)
    ax.plot (lon, lat, color= 'grey'      )
    ax.plot (lon[2500], lat[2500],  'gx' )
    ax.barbs(lon[::1000], lat[::1000], uwnd[::1000,400], vwnd[::1000,400], barbcolor = 'C0')
    ax.set_title(flightDate+'\nTrack')
    ax.grid(True)
    ax= fig.add_subplot(223)
    ax.plot(secs, theta[:,400],        'C4.',markersize = 1, label = 'WIND DIR rel to HEAD')
    ax.plot(secs, cross[:,400], color = 'C1', label = 'CROSS WIND (m/s)')
    ax.plot(secs, along[:,400], color = 'C2', label = 'ALONG WIND (m/s)')
    ax.xaxis.set_major_formatter(date_formatter)
    ax.set_xlabel('Time')
    ax.set_ylabel('Track-Relative')
    ax.grid(True)
    plt.legend(prop={'size': 6})
    ax= fig.add_subplot(224)
    ax.plot (time[::1000],time[::1000], color='grey')
    ax.barbs(time[::1000],time[::1000],   cross[::1000,400], 0*along[::1000,400], barbcolor='C1', alpha = 0.5)
    ax.barbs(time[::1000],time[::1000], 0*cross[::1000,400],   along[::1000,400], barbcolor='C2',  alpha = 0.5)
    ax.barbs(time[::1000],time[::1000],   cross[::1000,400],   along[::1000,400], barbcolor='C4')
    plt.figtext(0.64, 0.12, 'start')
    plt.figtext(0.64, 0.45, 'end')
    ax.set_xlabel('Time')
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    ax.grid(True)
    plt.savefig('./plots/' + radName + '_' + flightDate + 'flightpath_hrrr.png')
#    plt.show()
    '''

    return(cross, along, vel_corr_cross, vel_corr_along)