from io import StringIO
from matplotlib import pyplot as plt
import metpy.calc as mpcalc
from metpy.units import units
from metpy import plots
import pandas as pd
import numpy as np

def readSharppy(fileName):
    with open(fileName, 'r') as f:
        text = f.readlines()
    rest = ''.join(text[3:])
    data = ''.join(rest.split('%RAW%')[1]).split('%END%')[0].replace(' ', '')
    data = pd.read_csv(StringIO(data), sep=',', header=None, names=['LEVEL', 'HGHT', 'TEMP', 'DWPT', 'WDIR', 'WSPD'])
    data = data.replace(-9999, np.nan)
    data = data.loc[data['LEVEL'] >= 10]
    return data


fig = plt.figure()
skew = plots.SkewT(fig)
data = readSharppy('KCLL_1835Z_20221104_RAOB_SHARPPY.txt')
pressure = data['LEVEL'].values * units.hPa
temperature = data['TEMP'].values * units.degC
dewpoint = data['DWPT'].values * units.degC
virt = mpcalc.virtual_temperature_from_dewpoint(pressure, temperature, dewpoint)
profile_temp = mpcalc.parcel_profile(pressure, temperature[0], dewpoint[0])
profile_init_virtT = mpcalc.parcel_profile(pressure, virt[0], dewpoint[0])
profile_correct_virtT = mpcalc.parcel_profile(pressure, temperature[0], dewpoint[0], return_virt=True)

px = 1/plt.rcParams['figure.dpi']  
fig.set_size_inches(1920*px, 1080*px)
skew.plot_moist_adiabats(linewidths=0.5)
skew.plot(pressure, temperature, 'red', linewidth=1)
skew.plot(pressure, dewpoint, 'lime', linewidth=1)
skew.plot(pressure, virt, 'red', linestyle='dotted', linewidth=1)
fig.savefig("justthesounding.png")
skew.plot(pressure, profile_temp, 'black', linestyle='dashed', label='return_virt=False')
skew.plot(pressure, profile_correct_virtT, 'black', linestyle='solid', label='return_virt=True')
skew.ax.legend()
fig.savefig("soundingwithprofiles.png")
skew.ax.set_xlim()

plt.show()