import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
#import string as str
import datetime as dt
from scipy.interpolate import interp1d
from matplotlib import cm

def enfloat(strng,missingval=np.nan):
    """
    Converts a string to a float. Returns missingval if it can't.
    """
    try:
        return float(strng)
    except ValueError:
        return missingval

def get_line(file):
    line = ''
    while(str.strip(line) == ''):
        line = file.readline()
    return line

def read_sounding(line):
    """
    Reads one sounding from the U Wyoming text format and appends data to
    global variables.
    """
    global ipmax
    # the header line could be anything, assume we're at the first line of a record when called:
    #'91408 PTRO Koror, Palau Is Observations at 00Z 01 Aug 2017\n'
    word=str.strip(line).split()
    # find positions of spaces
    p = [ pos for pos, char in enumerate(line) if char==" "]
    # station identifiers
    station_num_id[irec]=int(word[0])
    station_wban_id.append(word[1])
    station_name.append(line[p[1]+1:p[4]])
    # timestamp
    #datestr=line[p[6]+1:p[7]-1]+line[p[7]::]
    datestr=word[-4][0:2]+' '+word[-3]+' '+word[-2]+' '+word[-1]
    #print(datestr)
    time.append(dt.datetime.strptime(datestr, '%H %d %b %Y'))
    
    # data header, always assumed in this form:
    #file.readline() #'\n' # get_line skips blank lines already
    get_line( file ) #'-----------------------------------------------------------------------------\n'
    get_line( file ) #'   PRES   HGHT   TEMP   DWPT   RELH   MIXR   DRCT   SKNT   THTA   THTE   THTV\n'
    get_line( file ) #'    hPa     m      C      C      %    g/kg    deg   knot     K      K      K \n'
    get_line( file ) #'-----------------------------------------------------------------------------\n'
    # now 10-column positional data table starts
    line = get_line( file )
    ip=0 # loop though vertical coordinate
    while not(str.strip(line).startswith('Station information and sounding indices')):
        #print(line)
        pres[irec,ip] = enfloat(line[0:7])
        hght[irec,ip] = enfloat(line[9:14])
        temp[irec,ip] = enfloat(line[16:21])
        dwpt[irec,ip] = enfloat(line[23:28])
        relh[irec,ip] = enfloat(line[31:35])
        mixr[irec,ip] = enfloat(line[37:42])
        drct[irec,ip] = enfloat(line[46:49])
        sknt[irec,ip] = enfloat(line[52:56])
        thta[irec,ip] = enfloat(line[58:63])
        thte[irec,ip] = enfloat(line[65:70])
        thtv[irec,ip] = enfloat(line[72:77])
        ip += 1
        line = get_line( file )
    ipmax = max(ipmax,ip)
    # skip summary data to end of record
    while not(str.strip(line).startswith('Precipitable water [mm] for entire sounding')):
        line = get_line( file )

#fileyears=range(2018,2018+1)
#filemonths=range(5,8+1)
fileyears=range(2018,2019) #excludes stop
filemonths=range(7,10)
nrec = 2 * 184 * len(fileyears) # number of sondes, assumes May-Oct inclusive each year, 2 sondes per day
npres=350
ipmax=0

# initialize variables (in global scope)
# housekeeping variables for each sonde
station_num_id=np.zeros(nrec,dtype=int)
station_wban_id=[] # empty list, append with station_wban_id.append("WBAN"). Lists are pythonic.
station_name=[]
time=[]
# sounding time,height variables
pres = np.nan+np.zeros((nrec,npres))
hght = np.nan+np.zeros((nrec,npres))
temp = np.nan+np.zeros((nrec,npres))
dwpt = np.nan+np.zeros((nrec,npres))
relh = np.nan+np.zeros((nrec,npres))
mixr = np.nan+np.zeros((nrec,npres))
drct = np.nan+np.zeros((nrec,npres))
sknt = np.nan+np.zeros((nrec,npres))
thta = np.nan+np.zeros((nrec,npres))
thte = np.nan+np.zeros((nrec,npres))
thtv = np.nan+np.zeros((nrec,npres))

# start main loop
irec=0 # initialize global sounding record index, 0-based
# loop over all files
for fy in fileyears:
    for fm in filemonths:
        filename="{:0>4d}{:0>2d}.txt".format(fy,fm)
        print(filename)
        with open(filename,'r') as file:
            # loop over each sounding
            line = get_line( file )
            while str.strip(line).startswith('University of Wyoming'):
                line = get_line( file )
            while not(str.strip(line).startswith(('Description of the data columns or sounding indices.','Description of the'))):
                read_sounding(line) # starts with the present line and reads more lines, appends data to sounding variables
                irec += 1
                line = get_line( file )

# truncate data
station_num_id=station_num_id[:irec]
pres=pres[:irec,:ipmax]
hght=hght[:irec,:ipmax]
temp=temp[:irec,:ipmax]
dwpt=dwpt[:irec,:ipmax]
relh=relh[:irec,:ipmax]
mixr=mixr[:irec,:ipmax]
drct=drct[:irec,:ipmax]
sknt=sknt[:irec,:ipmax]
thta=thta[:irec,:ipmax]
thte=thte[:irec,:ipmax]
thtv=thtv[:irec,:ipmax]

wspd = sknt*.5144
u = wspd*np.cos(np.radians(270-drct))
v = wspd*np.sin(np.radians(270-drct))
uanom = u - np.nanmean(u,axis=0)
vanom = v - np.nanmean(v,axis=0)
rhanom = relh - np.nanmean(relh,axis=0)
#Variable "time" tells us the datetime.

#Interpolate uanom, vanom, and rhanom to standard pressure levels for plotting. 
interp_levels = np.arange(1000,175,-25)
uanom_I = np.zeros([uanom.shape[0],interp_levels.size])
vanom_I = np.zeros([vanom.shape[0],interp_levels.size])
rhanom_I = np.zeros([rhanom.shape[0],interp_levels.size])
for i in np.arange(uanom.shape[0]):
  f = interp1d(pres[i,:],uanom[i,:])
  f2 = interp1d(pres[i,:],vanom[i,:])
  f3 = interp1d(pres[i,:],rhanom[i,:])
  uanom_I[i,:] = f(interp_levels)
  vanom_I[i,:] = f2(interp_levels)
  rhanom_I[i,:] = f3(interp_levels)
 
#Plot stuff.
fig,ax = plt.subplots(1,1,figsize=(10,4))
plt.contourf(np.arange(122),interp_levels,np.transpose(uanom_I),np.arange(-10,12,2),cmap=cm.get_cmap('RdBu'),extend="both")
ax.invert_yaxis()
ax.set_ylabel('Pressure (hPa)')
ax.set_xlabel('Date')
ax.set_xticks([0,17,36,55,71,90,109])
ax.set_xticklabels(['01-Jul','10-Jul','20-Jul','31-Jul','10-Aug','20-Aug','30-Aug'])
cb = plt.colorbar()
cb.set_label(r'Zonal wind anomaly (m $s^{-1}$)')
fig.savefig('Mactan_uanom.eps',dpi=300)

fig2,ax2 = plt.subplots(1,1,figsize=(10,4))
plt.contourf(np.arange(122),interp_levels,np.transpose(vanom_I),np.arange(-10,12,2),cmap=cm.get_cmap('RdBu'),extend="both")
ax2.invert_yaxis()
ax2.set_ylabel('Pressure (hPa)')
ax2.set_xlabel('Date')
ax2.set_xticks([0,17,36,55,71,90,109])
ax2.set_xticklabels(['01-Jul','10-Jul','20-Jul','31-Jul','10-Aug','20-Aug','30-Aug'])
cb2 = plt.colorbar()
cb2.set_label(r'Meridional wind anomaly (m $s^{-1}$)')
fig2.savefig('Mactan_vanom.eps',dpi=300)

fig3,ax3 = plt.subplots(1,1,figsize=(10,4))
plt.contourf(np.arange(122),interp_levels,np.transpose(rhanom_I),np.arange(-30,33,3),cmap=cm.get_cmap('RdBu'),extend="both")
ax3.invert_yaxis()
ax3.set_ylabel('Pressure (hPa)')
ax3.set_xlabel('Date')
ax3.set_xticks([0,17,36,55,71,90,109])
ax3.set_xticklabels(['01-Jul','10-Jul','20-Jul','31-Jul','10-Aug','20-Aug','30-Aug'])
cb3 = plt.colorbar()
cb3.set_label(r'Relative humidity anomaly (m $s^{-1}$)')
fig3.savefig('Mactan_rhanom.eps',dpi=300)