#%% 0. Imports
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import metpy.calc as mpcalc
from metpy.cbook import get_test_data
from metpy.plots import SkewT
from metpy.units import units


#%% 1. Import data (Não se preocupem muito com esta célula)
df = pd.read_fwf(get_test_data('may4_sounding.txt', as_file_obj=False),
                 skiprows=5, usecols=[0, 1, 2, 3, 6, 7],
                 names=['pressure', 'height', 'temperature', 'dewpoint', 'direction', 'speed'])

# Drop any rows with all NaN values for T, Td, winds
df = df.dropna(subset=('temperature', 'dewpoint', 'direction', 'speed'), how='all'
               ).reset_index(drop=True)


#%% 2. Retirar variáveis do ficheiro, atenção às unidades!
p = df['pressure'].values * units.hPa      # Pressão
T = df['temperature'].values * units.degC  # Temperatura
Td = df['dewpoint'].values * units.degC    # Temperatura do ponto de orvalho


#%% 3. Desenhar o tefigrama
#### Ajustar e inicializar tefigrama
plt.rcParams['figure.figsize'] = (11, 8)
skew = SkewT(rotation=45)

# Adicionar eixos relevantes
skew.plot_dry_adiabats(label='Dry adiabat',colors='grey',alpha=0.9)             # Adiabáticas secas
skew.plot_moist_adiabats(label='Moist adiabat',colors='red',alpha=0.9)          # Adiabáticas húmidas
skew.plot_mixing_lines(pressure=np.linspace(1000,250,1000) * units("hPa"),      # Razão de mistura
                       label='Mixing ratio', colors='tab:cyan')

# Nomes dos eixos
skew.ax.set_xlabel('Temperature (\N{DEGREE CELSIUS})', fontsize=18)
skew.ax.set_ylabel('Pressure (hPa)', fontsize=18)

# Ajuste dos eixos
skew.ax.set_ylim(1000, 250)
skew.ax.set_xlim(-10, 40)
skew.ax.tick_params(axis='both', which='major', labelsize=16)

# Desenhar as linhas de dados
skew.plot(p, T, 'r')
skew.plot(p, Td, 'g')

# Calculate LCL height and plot as black dot.
lcl_pressure, lcl_temperature = mpcalc.lcl(p[0], T[0], Td[0])
skew.plot(lcl_pressure, lcl_temperature, 'ko', markerfacecolor='black')

# Calculate full parcel profile and add to plot as black line
prof = mpcalc.parcel_profile(p, T[0], Td[0]).to('degC')
skew.plot(p, prof, 'k', linewidth=2)

# Shade areas of CAPE and compute
skew.shade_cape(p, T, prof)
cape=mpcalc.cape_cin(p,T,Td,prof)