import os
import glob
import numpy as np
from math import sqrt, log
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from collections import OrderedDict

linestyles = OrderedDict(
    [('solid',               (0, ())),
     ('loosely dotted',      (0, (1, 10))),
     ('dotted',              (0, (1, 5))),
     ('densely dotted',      (0, (1, 1))),

     ('loosely dashed',      (0, (5, 10))),
     ('dashed',              (0, (5, 5))),
     ('densely dashed',      (0, (5, 1))),

     ('loosely dashdotted',  (0, (3, 10, 1, 10))),
     ('dashdotted',          (0, (3, 5, 1, 5))),
     ('densely dashdotted',  (0, (3, 1, 1, 1))),

     ('loosely dashdotdotted', (0, (3, 10, 1, 10, 1, 10))),
     ('dashdotdotted',         (0, (3, 5, 1, 5, 1, 5))),
     ('densely dashdotdotted', (0, (3, 1, 1, 1, 1, 1)))])

def ustar(uref, zref, z0):
    """Calculate friction velocity for a neutral ABL logartithmic profile."""
    return KARMAN * uref / log((zref + z0) / z0)


def u_z(z, z0, zref, uref):
    frac = ustar(uref, zref, z0)/ KARMAN
    _ = frac * np.log(z/ z0)
    return _.reshape(-1, 1)
UREF = 3
ZREF = 5
KARMAN = 0.4
Cmu = 0.09
roughnessList = [0.001, 0.03,0.10, 0.25, 0.50, 1, 2]

Z = np.arange(1, 100, 2)
U_Z = np.zeros((Z.shape[0],len(roughnessList)))
for i, Z0 in enumerate(roughnessList):
    U_Z[:, i] = u_z(Z, Z0, ZREF, UREF)[:, 0] 
U_Z.shape, Z.shape   
((50, 7), (50,))
linestyles.keys()
odict_keys(['solid', 'loosely dotted', 'dotted', 'densely dotted', 'loosely dashed', 'dashed', 'densely dashed', 'loosely dashdotted', 'dashdotted', 'densely dashdotted', 'loosely dashdotdotted', 'dashdotdotted', 'densely dashdotdotted'])
fig, (ax1) = plt.subplots(1, 1, figsize=(10,10))
ax1.grid(True, which = "major", axis = "both")
ln_styles = ['solid', 'loosely dotted', 'dotted', 'densely dotted', 'loosely dashed',\
             'dashed', 'densely dashed','loosely dashdotted', 'dashdotted', 'densely dashdotted',\
             'loosely dashdotdotted', 'dashdotdotted','densely dashdotdotted']
for j in range(U_Z.shape[1]):
    ax1.plot(U_Z[:, j], Z, label='z0 = '+str(roughnessList[j])+' m', color = "k", ls = linestyles[ln_styles[j]])
    # add an ellipse
patches = []
ellipse = mpatches.Rectangle((-2, 0), 2, 4,color = 'r')
#patches.append(ellipse)
ax1.add_patch(ellipse)
ax1.legend(loc='upper left', shadow=True)   
ax1.set_ylabel('Height above ground [m]')
ax1.set_title("Wind speed [m/s]")
ax1.set_ylim((0,100))
ax1.set_xlim((-2,10))
ax1.annotate('Model breaks for high roughness', xy=(-1.0, 3.5), xytext=(-1.0, 15),
            arrowprops=dict(facecolor='red', edgecolor='red', shrink=0.05), color = 'red')
#plt.show()
plt.savefig("vel_Uz.png", dpi = 300)