#--------Helper script.  Not for courseware--------
#
# This script is used to interpolate radiosonde data to
# a set of standard pressure levels. It is used to turn
# the raw CEPEX radiosonde data into a dataset on a common
# set of levels, in order to allow the student to composite
# radiosondes without the distraction of interpolation.
# To use the script, edit the "datapath" string below to
# specify where the sonde data is located.  The script
# The script will put the output files in the directory
# it is run from, so you can move it to the directory where
# you want the processed data to appear before running it.
#---------------------------------------------------
#
#Currently, this script requires cdms to do the interpolation
#However,we need a general interpolation routine
#for using with numerical integrations of tabulated
#data, so maybe this could be re-written to use
#that instead.

#Edit datapath according to the local system
datapath = '/Users/rtp1/Havsornen/Teaching/GeoSci232/WorkbookDatasets/Chapter2Data/cepex_sondes/rawData/'


import cdms,MV,Numeric
from ClimateUtilities import *
from glob import glob
import os # Used to split the filename off from the path

#Function to embed a 1D array in a lev-lat-lon array
#  (necessary because the cdms pressure regrid
#   method only works on 3D arrays)/ data is a 1D
#   array, and levdata is the corresponding set of levels
def makeField3D(data,levdata):
    dat1 = MV.zeros((len(data),1,1),MV.Float)
    dat1[:,0,0] = data[:]
    # Make and install the level axis
    lev = cdms.createAxis(levdata,id='Level')
    lev.designateLevel()
    dat1.setAxis(0,lev)
    # Make and install the dummy lat and lon axes
    lat = cdms.createAxis([0.],id='Lat')
    lat.designateLatitude()
    dat1.setAxis(1,lat)
    lon = cdms.createAxis([0.],id='Lon')
    lon.designateLongitude()
    dat1.setAxis(2,lon)
    return dat1

#Read in data 
#   Note that cdms uses logarithmic interpolation for
#   pressure regridding, so we can't use zero pressure ever
#   as a coordinate.
#
cepex_header = ['time_min', 'time_sec', 'z', 'p', 'T', 'rh', 'Tdew', 'WindDir', 'WindSpeed']

#Loop over the filenames to process
for filename in glob(datapath+'*.dat'):
    print "Processing",filename
    #Get the lat/lon and time info from the header file
    infoFilename = filename.split('.')[0]+'.apa'
    f = open(infoFilename)
    header = f.readlines()
    f.close()
    #Only the first two lines of the cepex info file are of much interest.
    #Concatenate them into a description string
    description = header[0]+header[1]
    c = readTable(filename,cepex_header)

    p = Numeric.array(c['p'])
    T = Numeric.array(c['T']) + 273.15
    rh = Numeric.array(c['rh'])/100.

    #Make axis of standard levels
    standardLevels = cdms.createAxis([1000.- 20.*i for i in range(49)],id='Pressure')
    standardLevels.designateLevel()

    #Do the vertical intepolation
    T1 = makeField3D(T,p)
    T2 = T1.pressureRegrid(standardLevels)[:,0,0]
    rh1 = makeField3D(rh,p)
    rh2 = rh1.pressureRegrid(standardLevels)[:,0,0]

    #Install the desired data subset in a curve and dump to output file
    c1 = Curve()
    c1.addCurve(standardLevels,'p')
    c1.addCurve(T2,'T')
    c1.addCurve(rh2,'rh','Relative Humidity Fraction')
    c1.description = description
    #Split off the filename from the directory path, so the output
    #gets written to the current directory
    outfile = os.path.split(filename)[-1]
    c1.dump(outfile.split('.')[0]+'.txt')