# -*- coding: utf-8 -*-
"""
Created on Fri Nov 27 12:39:28 2020

@author: Sreenath
"""
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import pandas as pd
import os
from datetime import datetime
#import seaborn as sns
import matplotlib as mpl
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,
                               AutoMinorLocator)
#%% code wip

def rho(T, p):
    
    """
    
    Calculates air density (rho)
    
    """
    
    
    Rd = 287.0

#    Tv   = T * (1+0.61*qv) # Virtual temperature

    rho = p / (Rd * T) # Air density [kg m^-3]

    return(rho)

os.chdir(r'C:\Users\Sreenath\Documents\palm\Cheyenne\ches_forcings\HRRRB\hrrr\20190924\data')



check1 = '20190924_hrrr.t05z.wrfprsf00.small.nc'
check2 = '20190924_hrrr.t06z.wrfprsf00.small.nc'
check3 = '20190924_hrrr.t07z.wrfprsf00.small.nc'

#setup the xtime and pres variables
ds1_check = xr.open_dataset(check1)
hour = np.asarray([float(pd.to_datetime(ds1_check.initial_time0_hours[0].values).hour*60)])
initial_time0_hours = ds1_check['initial_time0_hours'].values
XTIME = xr.DataArray(hour, coords={'initial_time0_hours': initial_time0_hours},
             dims=['initial_time0_hours'])
ds1_check['XTIME'] = XTIME
ds1_check['XTIME'].attrs['long_name'] = 'minutes from forcings start'
ds1_check['XTIME'].attrs['units'] = ''

lv_ISBL0 = ds1_check.lv_ISBL0.values 
dummy = np.empty((40,110,111))

for i in range(np.shape(dummy)[1]):
    for j in range(np.shape(dummy)[2]):
        dummy[:,i,j] = lv_ISBL0

gridlat_0 = ds1_check.gridlat_0.values
gridlon_0 = ds1_check.gridlon_0.values

Pres = xr.DataArray(dummy, coords={'lv_ISBL0': lv_ISBL0,
                                       'gridlat_0':(['ygrid_0','xgrid_0'],gridlat_0),
                                       'gridlon_0':(['ygrid_0','xgrid_0'],gridlon_0)},
             dims=['lv_ISBL0','ygrid_0','xgrid_0'])
del dummy

ds1_check['Pres'] = Pres
ds1_check.Pres.attrs = ds1_check['Pres'].attrs


#rotate wind
ds1_check['Vearth'] = ds1_check.VGRD_P0_L100_GLC0*np.cos(ds1_check.gridrot_0) - \
                      ds1_check.UGRD_P0_L100_GLC0*np.sin(ds1_check.gridrot_0)
ds1_check['Uearth'] = ds1_check.VGRD_P0_L100_GLC0*np.sin(ds1_check.gridrot_0) + \
                      ds1_check.UGRD_P0_L100_GLC0*np.cos(ds1_check.gridrot_0)
#convert to potential temp
P0 = 100000 #Pa
ds1_check['Pt'] = ds1_check.TMP_P0_L100_GLC0*((P0/ds1_check.Pres)**(0.286))

#convert to vertical velocity
g = 9.80665 #m/s**2
ds1_check['WVEL'] =  -(ds1_check['VVEL_P0_L100_GLC0']/ (rho(ds1_check.TMP_P0_L100_GLC0,ds1_check.Pres)*g))


ds2_check = xr.open_dataset(check2)
hour = np.asarray([float(pd.to_datetime(ds2_check.initial_time0_hours[0].values).hour*60)])
initial_time0_hours = ds2_check['initial_time0_hours'].values
XTIME = xr.DataArray(hour, coords={'initial_time0_hours': initial_time0_hours},
             dims=['initial_time0_hours'])
ds2_check['XTIME'] = XTIME
ds2_check['XTIME'].attrs['long_name'] = 'minutes from forcings start'
ds2_check['XTIME'].attrs['units'] = ''

lv_ISBL0 = ds2_check.lv_ISBL0.values 
dummy = np.empty((40,110,111))

for i in range(np.shape(dummy)[1]):
    for j in range(np.shape(dummy)[2]):
        dummy[:,i,j] = lv_ISBL0

gridlat_0 = ds2_check.gridlat_0.values
gridlon_0 = ds2_check.gridlon_0.values

Pres = xr.DataArray(dummy, coords={'lv_ISBL0': lv_ISBL0,
                                       'gridlat_0':(['ygrid_0','xgrid_0'],gridlat_0),
                                       'gridlon_0':(['ygrid_0','xgrid_0'],gridlon_0)},
             dims=['lv_ISBL0','ygrid_0','xgrid_0'])
del dummy

ds2_check['Pres'] = Pres
ds2_check.Pres.attrs = ds1_check['Pres'].attrs

#rotate wind
ds2_check['Vearth'] = ds2_check.VGRD_P0_L100_GLC0*np.cos(ds2_check.gridrot_0) - \
                      ds2_check.UGRD_P0_L100_GLC0*np.sin(ds2_check.gridrot_0)
ds2_check['Uearth'] = ds2_check.VGRD_P0_L100_GLC0*np.sin(ds2_check.gridrot_0) + \
                      ds2_check.UGRD_P0_L100_GLC0*np.cos(ds2_check.gridrot_0)
#convert to potential temp
P0 = 100000 #Pa
ds2_check['Pt'] = ds2_check.TMP_P0_L100_GLC0*((P0/ds2_check.Pres)**(0.286))

#convert to vertical velocity
g = 9.80665 #m/s**2
ds2_check['WVEL'] =  -(ds2_check['VVEL_P0_L100_GLC0']/ (rho(ds2_check.TMP_P0_L100_GLC0,ds2_check.Pres)*g))




ds3_check = xr.open_dataset(check3)
hour = np.asarray([float(pd.to_datetime(ds3_check.initial_time0_hours[0].values).hour*60)])
initial_time0_hours = ds3_check['initial_time0_hours'].values
XTIME = xr.DataArray(hour, coords={'initial_time0_hours': initial_time0_hours},
             dims=['initial_time0_hours'])
ds3_check['XTIME'] = XTIME
ds3_check['XTIME'].attrs['long_name'] = 'minutes from forcings start'
ds3_check['XTIME'].attrs['units'] = ''

lv_ISBL0 = ds3_check.lv_ISBL0.values 
dummy = np.empty((40,110,111))

for i in range(np.shape(dummy)[1]):
    for j in range(np.shape(dummy)[2]):
        dummy[:,i,j] = lv_ISBL0

gridlat_0 = ds3_check.gridlat_0.values
gridlon_0 = ds3_check.gridlon_0.values

Pres = xr.DataArray(dummy, coords={'lv_ISBL0': lv_ISBL0,
                                       'gridlat_0':(['ygrid_0','xgrid_0'],gridlat_0),
                                       'gridlon_0':(['ygrid_0','xgrid_0'],gridlon_0)},
             dims=['lv_ISBL0','ygrid_0','xgrid_0'])
del dummy

ds3_check['Pres'] = Pres
ds3_check.Pres.attrs = ds1_check['Pres'].attrs

#rotate wind
ds3_check['Vearth'] = ds3_check.VGRD_P0_L100_GLC0*np.cos(ds3_check.gridrot_0) - \
                      ds3_check.UGRD_P0_L100_GLC0*np.sin(ds3_check.gridrot_0)
ds3_check['Uearth'] = ds3_check.VGRD_P0_L100_GLC0*np.sin(ds3_check.gridrot_0) + \
                      ds3_check.UGRD_P0_L100_GLC0*np.cos(ds3_check.gridrot_0)
#convert to potential temp
P0 = 100000 #Pa
ds3_check['Pt'] = ds3_check.TMP_P0_L100_GLC0*((P0/ds3_check.Pres)**(0.286))

#convert to vertical velocity
g = 9.80665 #m/s**2
ds3_check['WVEL'] =  -(ds3_check['VVEL_P0_L100_GLC0']/ (rho(ds3_check.TMP_P0_L100_GLC0,ds3_check.Pres)*g))





#ds3_check.SPFH_P0_L100_GLC0
#ds1_check.DSWRF_P0_L1_GLC0
#ds1_check.DLWRF_P0_L1_GLC0


#subset for only the needed data
check1 = ds1_check[['lv_DBLL7_l1','lv_DBLL7_l0','LAND_P0_L1_GLC0','TSOIL_P0_2L106_GLC0','HGT_P0_L100_GLC0',
             'XTIME', 'TMP_P0_L100_GLC0','Pres','SOILW_P0_2L106_GLC0','SPFH_P0_L100_GLC0',
             'Uearth','Vearth','WVEL','Pt','DSWRF_P0_L1_GLC0','DLWRF_P0_L1_GLC0']]
check2 = ds2_check[['lv_DBLL7_l1','lv_DBLL7_l0','LAND_P0_L1_GLC0','TSOIL_P0_2L106_GLC0','HGT_P0_L100_GLC0',
             'XTIME', 'TMP_P0_L100_GLC0','Pres','SOILW_P0_2L106_GLC0','SPFH_P0_L100_GLC0',
             'Uearth','Vearth','WVEL','Pt','DSWRF_P0_L1_GLC0','DLWRF_P0_L1_GLC0']]
check3 = ds3_check[['lv_DBLL7_l1','lv_DBLL7_l0','LAND_P0_L1_GLC0','TSOIL_P0_2L106_GLC0','HGT_P0_L100_GLC0',
             'XTIME', 'TMP_P0_L100_GLC0','Pres','SOILW_P0_2L106_GLC0','SPFH_P0_L100_GLC0',
             'Uearth','Vearth','WVEL','Pt','DSWRF_P0_L1_GLC0','DLWRF_P0_L1_GLC0']]

#check1.XTIME
"""subset for only the needed data

'ZS' and 'DZS', soil levels: 'lv_DBLL9_l1','lv_DBLL9_l0' 
they wll be used for depth of soil layer variable, cant find a thickness variable for 'DZS'

LANDMASK: LAND_P0_L1_GLC0

TMN - soil temperature at lower boundary: TSOIL_P0_2L106_GLC0,Soil temperature

H_stag =  HGT_P0_L105_GLC0
XTIME = XTIME, the variable we created

theta = TMP_P0_L100_GLC0, Temperature , pull in just temperature for now,
pres = PRES_P0_L105_GLC0
tk, temperature in kelvin = TMP_P0_L105_GLC0
TSLB, SOIL TEMPERATURE = TSOIL_P0_2L106_GLC0
SMOIS, SOIL MOISTURE m3 m-3 = SOILW_P0_2L106_GLC0
QVAPOR, Water vapor mixing ratio kg kg-1 = SPFH_P0_L100_GLC0, Specific humidity
U, x-wind = UGRD_P0_L105_GLC0
V, y-wind = VGRD_P0_L105_GLC0
W, z-wind = VVEL_P0_L105_GLC0, Vertical velocity (pressure) for now

"""


#concatenate on the time dimension
check_ds = xr.concat([check1, check2, check3], dim="initial_time0_hours")
check_ds = check_ds.sortby('lv_ISBL0',ascending=False).copy()

#check_ds.Pres

check_ds['Uearth']


#%%



# np.shape(check_ds.initial_time0_hours)

# check_ds.HGT_P0_L100_GLC0[0]
# check_ds.lv_ISBL0[39]

#flip the pressure coords

os.chdir(r'C:\Users\Sreenath\Documents\palm\Cheyenne\ches_forcings\WRF4PALM-master\wrf_output')

check_ds.to_netcdf('hrrr_check.nc')

#check_ds.close()
# check_ds = ds_hrrr.copy()