# Figure 6 of paper "A locally smoothed terrain-following coordinate"
# submitted to JAMES
# Author: Stephanie Westerhuis
# Date: February 1, 2021

import matplotlib
import numpy as np
import xarray as xr
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import datetime as dt
import sys
from scipy import interpolate, stats, spatial
import math
from pathlib import Path
import cartopy
import matplotlib.image as mpimg
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import seaborn as sns

# vertical profiles of lwc and total humidity
# compare SLEVE (100) and different LOSVEC configurations:
#   0 (670), 10 (671), 50 (672), 100 (673) filter iterations

def get_latlon_kdtree(lats, lons):
    return spatial.KDTree(np.array(list(zip(
        lons.ravel(), lats.ravel()))))

def get_ind_closest_to_latlon(lat, lon, lats, lons):

    kdtree = get_latlon_kdtree(lats, lons)
    x_shp = lons.shape[0]
    y_shp = lons.shape[1]
    data_index = np.arange(0, x_shp*y_shp).reshape(x_shp, y_shp)
    distance, i = kdtree.query((lon, lat))
    ind_lat_point, ind_lon_point = np.argwhere(data_index == i)[0]

    return ind_lat_point, ind_lon_point

def calc_alt(hhl):

    # altitude: height of full levels
    return hhl[1:,:,:]+(hhl[:-1,:,:]-hhl[1:,:,:])/2

# location payerne
(lat_pay, lon_pay) = (46.81303, 6.94302)

setups = ['sleve','losvec_0','losvec_10','losvec_50','losvec']
labels = ['SLEVE',
          'LOSVEC: N = 0',
          'LOSVEC: N = 10',
          'LOSVEC: N = 50',
          'LOSVEC: N = 100']
colors = ['yellowgreen',
          'orange',
          'magenta',
          'grey',
          'mediumblue']
linestyles = ['solid',
              (0, (1, 1)), # dotted
              (0, (3, 5, 1, 5)), #dash-dotted
              (0, (3, 5, 1, 5, 1, 5)), #dash-dot-dotted
              (0, (5, 5))] # dashed
              
fig, (ax1, ax2) = plt.subplots(figsize=(7.25,4), dpi=300,
                       nrows = 1, ncols = 2, sharey=True)

for ii, setup in enumerate(setups):
    f = f'data/fcst_{setup}/lfff00150000.nc'
    c = f'data/const_{setup}.nc'
    hhl = xr.open_dataset(c).isel(time=0).HHL.values
    alt = calc_alt(hhl)
    ds = xr.open_dataset(f).isel(time=0)
    try:
        lats = ds.lat.values
        lons = ds.lon.values
    except AttributeError:
        lats = ds.lat_1.values
        lons = ds.lon_1.values
    qc = ds.QC.values*1000  # liquid water content
    qt = ds.QC.values*1000+ds.QV.values*1000    # total humidity
    ilat,ilon = get_ind_closest_to_latlon(lat_pay,lon_pay,lats,lons) 

    # lwc
    ax1.plot(qc[:,ilat,ilon],
             alt[:,ilat,ilon],
             color = colors[ii],
             linestyle = linestyles[ii],
             marker = 'o',
             label = labels[ii],
             )

    # tot hum.
    ax2.plot(qt[:,ilat,ilon],
             alt[:,ilat,ilon],
             color = colors[ii],
             linestyle = linestyles[ii],
             marker = 'o',
             label = labels[ii],
             )

for ax in [ax1,ax2]:
    ax.set_ylim(500,800)
ax1.set_xlabel('Liquid water content [g/kg]')
ax1.set_ylabel('Elevation [m asl]')
ax1.legend(numpoints=2, handlelength=5)

ax2.set_xlabel('Total humidity [g/kg]')
ax2.set_xlim(2.5,3.05)

plt.tight_layout()

# save
out_name = '6_comp_setups'
fig.savefig(f'plots/{out_name}.png')