# coding: utf-8

# In[3]:

# ORIGINAL HEADER:
# Brody Fuchs, CSU, November 2015
# brfuchs@atmos.colostate.edu


## Update the skewT code to be smarter using classes and what not
# use the KUILsounding.txt as an example or KBMXsounding.txt

# Modified by Zane Martin, Columbia University, zkm2102@columbia.edu for PISTON cruise and for use with iPython 
# notebook

# Start out by just plotting the raw data first

import numpy as np
import matplotlib.pyplot as plt
from skewPy import SkewT
import os
from datetime import datetime
import glob
import sys

import argparse


# In[9]:

parser = argparse.ArgumentParser(description='Put in a file to be processed')

#parser.add_argument('--noarg', action="store_true", default=False)
parser.add_argument('--file', action="store", dest="file", default=None)

# location of sounding files
parser.add_argument('--filepath', action='store', dest='filepath', default='.')
# where you want to output skewTs
parser.add_argument('--outpath', action='store', dest='outpath', default='.')
# the first part of the sounding raw data file, need it to search for possible files
                    # other possibilities are 'K' for K*** UWYO soundings, could also be 'EDT' for Vaisala soundings
parser.add_argument('--prefix', action="store", dest="prefix", default='')
# This can currently be 'EC' for Canadian format, 'UWYO' for Wyoming soundings or 'EDT' for Vaisala CSU soundings
# Or XML, so some dumb*** reason. Gave us something to do on the SPURS2 cruise, I suppose
parser.add_argument('--format', action='store', dest='format', default='XML')
# put this in if want to override the station name in the title of the skewT, default is fmt above
                                    # otherwise just put None
parser.add_argument('--station', action='store', dest='station', default=None)
# turning this switch on will just print more stuff, maybe needed if trying to debug
parser.add_argument('--debug', action='store', dest='debug', type=int, default=0)
# This will process a certain file in the filepath. The number will correspond to the index of the file
parser.add_argument('--number', action='store', dest='number', default=None)
# This will process the last file in the filepath
parser.add_argument('--last', action='store', dest='last', default=None)


pargs = parser.parse_args()



# In[2]:


##### PARAMETERS YOU CAN CHANGE HERE #########

file_path = './' 
out_path = './' 
prefix = 'Revelle' # the first part of the sounding raw data file, need it to search for possible files
#                     # other possibilities are 'K' for K*** UWYO soundings, could also be 'EDT' for Vaisala soundings
fmt = 'XML' 
# station_name = None # put this in if want to override the station name in the title of the skewT, default is fmt above
#                                     # otherwise just put None

########### DON'T NEED TO WORRY ABOUT ANYTHING BELOW HERE ############


if len(sys.argv) > 1:

    if pargs.file is not None:
    	sounding_files = [pargs.file]

    elif pargs.number is not None:
        number_loc = int(pargs.number) + 1
        sounding_files = glob.glob('%s/%s*'%(pargs.filepath, pargs.prefix))[number_loc]

    elif pargs.number is not None: 
        sounding_files = [glob.glob('%s/%s*'%(pargs.filepath, pargs.prefix))[-1]]
    else:
        sounding_files = glob.glob('%s/%s*'%(pargs.filepath, pargs.prefix))

else: 
    sounding_files = sorted(glob.glob('%s/%s*'%(pargs.filepath, pargs.prefix)))


print('processing the following files: {}', format(sounding_files))



CAPEs = [] ## added during SPURS-2 for CAPE v dBZ plots
p_top = []
datestr = []
TW = np.arange(0)
RH_sfc = []
q_sfc = []
q_ml = []
for fname in sounding_files:
    #file_title = os.path.basename(fname)[4:-4]

    #print 'Processing %s'%os.path.basename(fname)
    S = SkewT.Sounding(fname, fmt=pargs.format, station_name=pargs.station)
    CAPE, pres_last, date, tw, RH, q, qm = S.cape()
    datestr.append(date)
    CAPEs.append(CAPE) ## added during SPURS-2 for CAPE v dBZ plots
    p_top.append(pres_last)
    TW = np.append(TW,tw)
    RH_sfc.append(RH)
    q_sfc.append(q)
    q_ml.append(qm)

    plt.savefig('%s/%s_sounding.png'%(pargs.outpath, S.sounding_date), dpi=150)
    plt.close('all')

print('\n%d total soundings in %s\n', (len(sounding_files), pargs.filepath))
#np.savetxt('CAPE.txt', np.transpose([datestr, CAPEs, p_top]), delimiter=',', fmt=(('%s'),('%.2f'),('%.2f'))) ## added during SPURS-2 for CAPE v dBZ plots
np.savez_compressed('CAPE.npz', times=datestr, cape=CAPEs, p_top=p_top, TW=TW, RH_sfc=RH_sfc, q_sfc=q_sfc, q_ml=q_ml)





# In[14]:




# In[ ]: