#!/usr/bin/env python
"""
Utilities for reading, analyzing, and visualizing idealized or observed
atmospheric profiles/soundings. 

Adapated from the `pywrfplotutils` package by Geiur Arne Waagbm, 
http://code.google.com/p/pywrfplot

Author: Daniel Rothenberg (darothen@mit.edu)

"""

import math
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d

import matplotlib.pyplot as plt

## Settings controlling the extent of the skew-T plots generated by this
## package
P_top, P_bot = 1e4, 1e5 # Top/bottom pressure levels, Pa
T_base = 300.0
T_zero = 273.15

## Thermodynamic constants
L = 2.501e6 # Latent heat of vaporization for water
R = 287.04  # gas constant air
Rd = R
Rv = 461.5  # gas constant vapor
eps = R/Rv
cp = 1005.
cv = 718.
kappa = (cp-cv)/cp
g = 9.81
gamma_d = 9.8 # dry adiabatic lapse rate, C/km

#####################################
## Filtering and other functions
##

def filter_121(data, niter=1, retain_endpoints=True):
    """ Iteratively apply a 1-2-1 filter to a given sequence of data.

    Parameters
    ----------
    data : an array of the data to smooth
    niter (optional) : the number of times to apply the filter
    retain_endpoints (optional) : boolean indicating whether to keep end datapoints fixed

    Returns
    -------
    an array of the same shape as data which has been filtered

    """
    data_old = data.copy()
    
    data_iter = data.copy()
    for i in xrange(niter):
        data_iter[1:-1] = 0.25*(2.*data[1:-1] + data[:-2] + data[2:])
        if retain_endpoints:
            data_iter[0] = data[0]
            data_iter[-1] = data[-1]
        else:
            data_iter[0] = (2.*data[0] + data[1])/3. 
            data_iter[-1] = (2.*data[-1] + data[-2])/3.
        data[:] = data_iter[:]
        
    return data_iter

#####################################
## Auxiliary thermodynamic functions
##

def calc_dewpoint(T, P, r):
    """ Compute dewpoint from temperature, pressure, and water vapor, using the
    [Magnus formula](http://en.wikipedia.org/wiki/Dew_point)

    Parameters
    ----------
    T : temperature in degrees C
    P : Pressure in hPa
    r : water vapor mixing ratio in kg/kg

    Returns
    -------
    dewpoint temperature in degrees C

    """
    ## Determine relative humidity by calculating actual/supersaturation 
    ## water vapor pressures
    epsilon = 0.622
    es_T = calc_es(T)
    P = P*100.
    e_env = P*r/epsilon
    RH = (e_env/es_T)*100.
    
    ## Constants for the Magnus approximation, from Bolton (1980; MWR)
    a = 6.112 # mb
    b = 17.67 #
    c = 243.5 # C
    gamma = lambda T, RH: np.log(RH/100.) + b*T/(c+T)
    dewpoint = c*gamma(T, RH)/(b-gamma(T, RH))

    return dewpoint

def calc_es(T):
    """ Calculate the equilibrium saturation vapor pressure for a parcel of
    air with given temperature, following Rogers and Yau, Formula 2.17

    Parameters
    ----------
    T : temperature in degrees C

    Returns
    -------
    equilibrium saturation vapor pressure in Pa 

    """
    arg = 17.67*T/(T + 243.5)
    return 611.2*np.exp(arg)

def gamma_s(T, p):
    """ Calculates the moist adiabatic lapse rate at a given temperature
    and pressure, using Rogers and Yau, Formula 3.16

    .. note:: although typically given as :math:`\frac{dT}{dz}`, here we have
        used the dry adiabatic lapse rate and the hydrostatic equation to 
        compute instead :math:`\frac{dT}{dP}` 

    Parameters
    ----------
    T : temperature in degrees C
    P : pressure in Pa

    Returns
    -------
    moist adiabatic lapse rate, in deg C / Pa

    """
    a = 2./7.
    b = eps*L*L/(R*cp)
    c = a*L/R

    esat = calc_es(T)
    wsat = eps*esat/(p-esat) # Rogers and Yau 2.18
    numer = a*(T+T_zero) + c*wsat
    denom = p * (1. + b*wsat/((T+T_zero)**2))

    return numer/denom # Rogers and Yau 3.16

def calc_theta(T, P, r=0.):
    """ Compute potential temperature from ambient temperature, pressure, and moisture

    .. note:: If `r` is provided, then moist potential temperature will be estimated

    Parameters
    ----------
    T : temperature in degrees C
    P : pressure in hPa
    r : water vapor mixing ratio in kg/kg (optional)

    Returns
    -------
    potential temperature in K

    """
    T_k = T + 273.15
    if r <= 0.: ## calculate dry potential temperature
        theta = T_k*((1000./P)**(Rd/cp))
    else: ## else, account for moisture
        Rm = R*(1.0 + 0.608*r)
        cpm = cp*(1. + 0.887*r)
        cvm = cv*(1. + r)
        theta = T_k*((1000./P)**(Rm/cpm))

    return theta

def calc_theta_e(T, P, r):
    """ Compute equivalent potential temperature from ambient temperature, 
    pressure, and moisture. Uses the formula from Bolton, 1980 (MWR)

    Parameters
    ----------
    T : temperature in degrees C
    P : pressure in hPa
    r : water vapor mixing ratio in kg/kg

    Returns
    -------
    equivalent potential temperature in K

    """
    ## Calculate the LCL temperature
    T_d = calc_dewpoint(T, P, r)
    T_lcl = calc_T_lcl(T, T_d)

    ## Calculate theta_e
    theta_m = calc_theta(T, P, r)
    theta_e = theta_m*np.exp((3376./T_lcl - 2.54)*r*(1.+0.81*r))

    return theta_e

def calc_T_lcl(T, T_d):
    """ Compute the temperature at lifting condensation level given an initial 
    temperature and water vapor in terms of dewpoint. Uses the
    formula from Bolton, 1980 (MWR)

    Parameters
    ----------
    T : initial temperature in degrees C
    T_d : initial dewpoint temperature in degrees C

    Returns
    -------
    temperature at lifting condensation level in degrees ??

    """
    T = T + 273.15
    T_d = T_d + 273.15
    T_lcl = 56. + 1./(1./(T_d - 56.) + np.log(T/T_d)/800.)

    return T_lcl

#####################################
## Utilities for reading soundings
##

def read_sounding(filename):
    """ Read a sounding from the UWYO archive at 
    http://weather.uwyo.edu/upperair/sounding.html

    Parameters
    ----------
    filename : the location of the sounding file to read and import

    Returns
    -------
    A pd.DataFrame with columns containing all the data from the sounding
        file

    """
    with open(filename, 'r') as f:
        lines = f.readlines()
    # skip header
    lines = lines[4:]

    def _fix_string(s):
        s = s.strip()
        if s: 
            return float(s)
        else: 
            return np.NaN            

    data = []
    for line in lines:
        bits = line.split()
        if len(bits) < 11: 
            
            pressure = line[:7]
            height = line[9:14]
            temperature = line[16:21]
            dewpoint = line[23:28]
            rh = line[30:35]
            mixing_ratio = line[37:42]
            direction = line[44:49]
            sknt = line[51:56]
            theta = line[58:63]
            theta_e = line[64:70]
            theta_v = line[71:]

            proc_line = map(_fix_string,
                            [pressure, height, temperature, dewpoint, rh, mixing_ratio,
                             direction, sknt, theta, theta_e, theta_v])
        else:
            proc_line = map(_fix_string, bits)
        data.append(proc_line)

    columns = [
        "pressure",     # hPa
        "height",       # meters
        "temperature",  # C
        "dewpoint",     # C
        "rh",           # %
        "mixing_ratio", # g/kg
        "direction",    # degrees
        "sknt",         # knots
        "theta",        # K
        "theta_e",      # K
        "theta_v",      # K
    ]
    data_df = pd.DataFrame(data, columns=columns)
    return data_df

#####################################
## Plotting Functions
##

def plot_hodograph(sounding, figsize=(8,8), **kawrgs):
    """ Plot a hodograph generated from a given sounding.

    .. note:: Requires that the wind direction and speed be passed with the
        sounding; in the future, will be able to convert back from component
        winds.

    Parameters
    ----------
    sounding : pd.DataFrame
        Environmental sounding data following the convention defined in 
        `plot_skewt`.
    ax : axis
        If present, will draw the plot on the user-given axis
    figsize : tuple
        If present, will force the plot figure size

    Returns
    -------
    axis on which the diagram was plotted

    """    
    def _choose_kwargs(key, default):
        if key in kwargs:
            return kwargs[key]
        else: 
            return default

    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, polar=True)

    ## Select the data
    zs = sounding.height.values

    ## Convert the winds if necessary. We want the wind direction in degrees
    ## and the speed magnitude in knots.
    if (("sknt" in sounding) and ("direction" in sounding)):
        speeds = sounding.sknt.values
        dirs = sounding.direction.values
    elif (("us" in sounding) and ("vs" in sounding)):
        us, vs = sounding['us'].values, sounding['vs'].values
        raise ValueError("conversion from vector to absolute winds not implemented yet")
    else:
        raise ValueError("Wind data not provided with sounding.")
    dirs = dirs*np.pi/180. # convert to radians

    ax.plot(dirs, speeds, 'or', linewidth=2)
    ax.set_rmax(60.)
    #ax.set_rgrids([0., 20., 40., 60.,])

    return ax


def plot_skewt(sounding, ax=None, figsize=(8,8),
    lift_parcel=False, plot_winds=False, diags=[],
    **kwargs):
    """ Plot a Skew-T/Log-P diagram generated from a given sounding, and 
    perform any additional analytical/diagnostic calculations specified.

    This basic routine was adopted from the pywrfplot utility, particularly
    the implementation of skewed axes. Additional diagnostic calculations 
    have been added later.

    Parameters
    ----------
    sounding : pd.DataFrame 
        Environmental sounding data with the following columns:
        - *temperature* : environmental temperature in degrees C
        - *pressure* : environmental pressure in hPa
        - *dewpoint* : environmental dewpoint in degrees C

        **Necessary for ``lift_parcel == True``**
        - *height* : measurement altitude in meters
        - *mixing_ratio* : environmental water vapor mixing ratio in g/kg
        **Necessary for ``plot_winds == True``**
        - *us*, *vs* : zonal and meridional winds, knots *or*
        - *sknt*, *direction* : wind absolute speed, knots, and direction, degrees
    ax : axis
        If present, will draw the plot on the user-given axis
    figsize : tuple
        If present, will force the plot figure size
    lift_parcel : Boolean
        Perform a lifted-parcel calculation from the surface to plot on the
        diagram
    plot_winds : Boolean
        Plot wind barbs with height on the diagram
    diags : list
        Optional diagnostics to calculate for the sounding. Possible values
        include - [, ]

    Returns
    -------
    axis on which the diagram was plotted

    """
    def _choose_kwargs(key, default):
        if key in kwargs:
            return kwargs[key]
        else: 
            return default

    skewness = _choose_kwargs("skewness", 37.5)
    P_b = _choose_kwargs("P_b", 105000.) # Bottom pressure, Pa
    P_t = _choose_kwargs("P_t", 10000.) # Top pressure, Pa
    dp = _choose_kwargs("dp", 50.) # Pressure increment, Pa
    T_left = _choose_kwargs("T_left", -40.)
    T_right = _choose_kwargs("T_right", 50.)
    label_altitudes = _choose_kwargs("label_altitudes", False)
    title = _choose_kwargs("title", None)

    plevs = np.arange(P_b, P_t-1, -dp)

    ## Start making figure
    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111)
    ax.grid(None, which='both', axis='both', alpha=0.)

    def _skewnessTerm(P):
        return skewness * np.log(P_bot/P)

    def _isotherms():
        for temp in np.arange(-140,50,10):
            ax.semilogy(temp + _skewnessTerm(plevs), plevs,  basey=math.e, \
                         color = ('blue' if temp <= 0 else 'red'), \
                         linestyle=('solid' if temp == 0 else 'dashed'), linewidth = .5)

    def _isobars():
        for n in np.arange(P_bot,P_t-1,-10**4):
            ax.plot([-40,50], [n,n], color = 'black', linewidth = .5)

    def _dry_adiabats():
        for tk in T_zero+np.arange(-30,210,10):
            dry_adiabat = tk * (plevs/P_bot)**kappa - T_zero + _skewnessTerm(plevs)
            ax.semilogy(dry_adiabat, plevs, basey=math.e, color = 'brown', \
                         linestyle='dashed', linewidth = .5)

    def _moist_adiabats():
        ps = [p for p in plevs if p<=P_bot]
        for temp in np.concatenate((np.arange(-40.,10.1,5.),np.arange(12.5,45.1,2.5))):
            moist_adiabat = []
            for p in ps:
                temp -= dp*gamma_s(temp,p)
                moist_adiabat.append(temp + _skewnessTerm(p))
            ax.semilogy(moist_adiabat, ps, basey=math.e, color = 'green', \
                         linestyle = 'dotted', linewidth = .5)

    def _mix_rats():
        ## Following SkewT (https://github.com/tchubb/SkewT/blob/master/skewt/SkewT.py)
        rs = np.array([0.0001, 0.0004, 0.001, 0.002, 0.004, 0.007, 0.01, 0.016, 0.24, 0.032])
        plevs_trunc = plevs[plevs>=70000.]

        e = np.array([plevs_trunc*r/(0.622+r) for r in rs])
        t = 243.5/(17.67/np.log(e/611.2) - 1.)

        for tt, mr in zip(t, rs):
            ax.semilogy(tt + _skewnessTerm(plevs_trunc), plevs_trunc, basey=math.e, \
                        color = "green", linestyle="dashed", lw=1.)
            if (tt[0] > T_left) and (tt[-1] < T_right):
                if mr*1e3 < 1: fmt = "%4.1f"
                else: fmt = "%d"
                ax.text(tt[10]+2. + _skewnessTerm(plevs_trunc[10]), plevs_trunc[10], 
                        fmt % (mr*1e3), 
                        color="green", fontsize=8, ha="center", va="bottom")

    def _temperature(temperature, pressure):
        linestyle = _choose_kwargs("T_style", "-")
        color = _choose_kwargs("T_color", "black")
        s = linestyle+"k" # note that color will always be overridden
        ax.semilogy(temperature + _skewnessTerm(pressure), pressure, s, basey=math.e, color=color, \
                     linestyle='solid', linewidth = 2.)

    def _dewpoint(dewpoints, pressure):
        linestyle = _choose_kwargs("Td_style", "-")
        color = _choose_kwargs("Td_color", "blue")
        s = linestyle+"k" # note that color will always be overridden
        ax.semilogy(dewpoints + _skewnessTerm(pressure), pressure, s, basey=math.e, color=color, \
                     linestyle='solid', linewidth = 2.)

    _isotherms()
    _isobars()
    _dry_adiabats()
    _moist_adiabats()
    _mix_rats()

    pressure = sounding['pressure']*100. # hPa -> Pa
    height = sounding['height'] # m
    temperature = sounding['temperature'] # C
    dewpoint = sounding['dewpoint']

    _temperature(temperature, pressure)
    _dewpoint(dewpoint, pressure)

    if plot_winds:
        barb_color = _choose_kwargs("barb_color", 'k')
        T_location = _choose_kwargs("barb_loc", 45.)

        if (("us" in sounding) and ("vs" in sounding)):
            us, vs = sounding['us'].values, sounding['vs'].values
        elif (("sknt" in sounding) and ("direction" in sounding)):
            ## Calculate zonal/meridional wind vectors
            model_dir = sounding.direction
            model_wspd = sounding.sknt

            dirs = np.pi*model_dir/180. # in radians
            us = -1.*model_wspd*np.sin(dirs)
            vs = -1.*model_wspd*np.cos(dirs)

            us = us.values
            vs = vs.values

        wind_ax = ax.twinx()
        wind_ax.grid(None)
        for spine in wind_ax.spines.itervalues():
            spine.set_visible(False)
        wind_ax.xaxis.set_visible(False)
        wind_ax.yaxis.set_visible(False)

        wind_ax.barbs(np.ones_like(us)[::2]*T_location, pressure[::2], us[::2], vs[::2],
                      length=6, barbcolor=barb_color)

        wind_ax.get_shared_y_axes().join(ax, wind_ax)
        wind_ax.semilogy()

    if lift_parcel:
        first_index = _choose_kwargs("first_index", 1)
        parcel_sounding = compute_lifted_parcel(sounding, first_index)
        misc = parcel_sounding.misc
        color = _choose_kwargs("parcel_color", "red")

        p_pressure, p_temperature = parcel_sounding['pressure']*100., parcel_sounding['temperature']
        ax.semilogy(p_temperature + _skewnessTerm(p_pressure), p_pressure, basey=math.e, color=color,
                     linestyle='dashed', linewidth=1.5)

        parcel_sounding_iso = parcel_sounding.dropna()
        p_pres_iso, p_temp_iso = parcel_sounding_iso['pressure']*100., parcel_sounding_iso['t_rsfc_iso']
        ax.semilogy(p_temp_iso + _skewnessTerm(p_pres_iso), p_pres_iso, basey=math.e, color=color,
                     linestyle='dashed', linewidth=1.5)

        ax.plot(misc['T_lcl'] + _skewnessTerm(misc['P_lcl']*100.), misc['P_lcl']*100., 
                     color=color, marker='o', ms=4, zorder=-1)

        ## Perform additional diagnostics and print output
        print "sounding diagnostics"
        print "--"*35
        fmt = "    {key:s} {val:6.1f}"

        misc = parcel_sounding.misc
        print fmt.format(key="LCLT", val=misc["T_lcl"])
        print fmt.format(key="LCLP", val=misc["P_lcl"])

        if "CAPE" in diags:
            parcel_sounding = compute_CAPE(sounding, parcel_sounding)
            print fmt.format(key="CAPV", val=misc["CAPV"])
            print fmt.format(key="EQTV", val=misc["EL_p"])
            print fmt.format(key="LFCV", val=misc["LFC_p"])


    #axis([40, 50, P_b, P_t])
    ax.set_xlim(T_left, T_right)
    ax.set_ylim(P_b, P_t)

    ax.set_xlabel('Temperature ($^{\circ}$C)', fontsize=10, color='k')
    xtick = np.arange(-40, 51, 5)
    ax.set_xticks(xtick)
    ax.set_xticklabels(['' if tick%10 != 0 else str(tick) for tick in xtick])

    ax.set_ylabel('Pressure (hPa)', fontsize=10, color='k')
    ytick = np.arange(P_bot, P_t-1, -10**4)
    ax.set_yticks(ytick)
    ax.set_yticklabels(["%4d" % tick for tick in ytick/100.])

    ## Label pressure altitudes
    if label_altitudes:
        p_subset = pressure[::6]
        z_subset = height[::6]
        for p, z in zip(p_subset, z_subset):
            if p < P_t: continue
            ax.text(T_left+0.5 , p,
                    "{:>6.1f} m".format(z), fontsize=8, ha="left", va="center")

    ## Add title
    if title:
        plt.title(title, loc="left", fontsize=12)

    if lift_parcel:
        return ax, parcel_sounding
    else:
        return ax

def compute_lifted_parcel(sounding, height_cutoff=10000.0, first_index=1):
    """ Lift a parcel from a surface with respect to a given sounding

    Parameters
    ----------
    sounding : pd.DataFrame
        The environmmental profile, conforming to the requirements in ``plot_skewt``
    height_cutoff : float
        Altitude at which to stop lifting the parcel
    first_index : int
        Index of the level from which valid data begins; the user might need to specify this

    Returns
    -------
    A new pd.DataFrame with the profile of the lifted parcel ascent with respect to
    height, pressure, temperature, and virtual temperature

    """
    height_cutoff = 1e5
    dp = 1. # hPa

    trunc_sounding = sounding[sounding['height'] < height_cutoff].dropna()
    p_interp = interp1d(trunc_sounding.height, trunc_sounding.pressure, 'slinear')
    first = trunc_sounding.index[first_index]

    ## 1) Compute the surface dry adiabat and lift it all the way to the
    ##    top of the sounding. We'll use it to find the LCL and then replace
    ##    everything above it with the appropriate moist adiabat. 
    t0, p0, r0 = trunc_sounding[["temperature", "pressure", "mixing_ratio"]].ix[first]
    print "Lifting parcel from T=%3.1f C, P=%4d hPa, Qv=%2.1f g/kg" % (t0, p0, r0)
    p_all = trunc_sounding.pressure.values
    p_dry = np.arange(p_all[0], p_all[-1], -dp)

    # dry adiabat
    theta0 = (t0+273.15)*((1e3/p0)**(Rd/cp)) # Kelvin
    ts_from_theta0 = theta0*(p_all/1e3)**(Rd/cp) - 273.15 # deg C
    ts_theta0_interp = interp1d(p_all[::-1], ts_from_theta0[::-1], 'slinear')
    T_dry_adiabat = ts_theta0_interp(p_dry)

    # sfc mixing ratio isopleth
    r0 *= 1e-3
    e = (p_dry*100.)*r0/(0.622+r0)
    ts_from_r0 = 243.5/(17.67/np.log(e/611.2) - 1.)
    T_r_isopleth = ts_from_r0
    #print r0, ts_from_r0[0], e[0]

    # compute where the two sfc temperature curves intersect
    P_lcl = np.interp(0., T_r_isopleth - T_dry_adiabat, p_dry)
    T_lcl = np.interp(P_lcl, p_dry[::-1], T_dry_adiabat[::-1])

    #print P_lcl, T_lcl

    ## 2) Lift the parcel adiabatically from the LCL pressure level
    p_moist = np.arange(P_lcl, 100., -dp)
    T_moist_adiabat = np.zeros_like(p_moist)
    T_moist_adiabat[0] = T_lcl
    for i in xrange(1, len(T_moist_adiabat)):
        dp_layer = p_moist[i] - p_moist[i-1]
        dt_dp = gamma_s(T_moist_adiabat[i-1], 100.*(p_moist[i-1]+p_moist[i])/2.)
        T_moist_adiabat[i] = T_moist_adiabat[i-1] + 100.*dt_dp*dp_layer

    ## 3) Setup output

    # shrink down the dry-lifting arrays to only go up to the LCL
    T_dry_adiabat = T_dry_adiabat[p_dry > P_lcl]
    T_r_isopleth = T_r_isopleth[p_dry > P_lcl]
    p_dry = p_dry[p_dry > P_lcl]

    p_combined = np.concatenate((p_dry, p_moist))
    t_combined = np.concatenate((T_dry_adiabat, T_moist_adiabat))
    t_iso = np.empty(t_combined.shape)
    t_iso[p_combined <= P_lcl] = np.nan
    t_iso[p_combined > P_lcl] = T_r_isopleth

    output = pd.DataFrame({'pressure': p_combined, 'temperature': t_combined, 
                           't_rsfc_iso': t_iso, })
    output.misc = { 'T_lcl': T_lcl, 'P_lcl': P_lcl, }

    return output

def compute_CAPE(sounding, parcel_sounding=None, lift_parcel=False):
    """ Perform a simple estimate of the profile convective available potential
    energy, accounting for moisture by using virtual temperature. The current
    scheme will crudely assume that, starting from the top of the atmosphere,
    the first time parcel and environment share the same temperature, this is the
    equilibrium level, and the next time they are equal this is the level of
    free convection.

    Parameters
    ----------
    sounding : pd.DataFrame
        An environmental sounding
    parcel_sounding : pd.DataFrame
        An ascending parcel's sounding, or
    lift_parcel : Boolean
        If True, will automatically recompute the lifted parcel.

    Returns
    -------
    Modified `parcel_sounding` containing updated entries in the `misc` dictionary
    corresponding to CAPE and the EL/LFC

    """

    from scipy.integrate import trapz

    if not lift_parcel:
        assert parcel_sounding is not None
    else:
        parcel_sounding = compute_lifted_parcel(sounding)

    first, last = sounding.index[1], sounding.index[-1]
    h_start, h_end = sounding.height.ix[first], sounding.height.ix[last]
    dz = 1. # meters
    z_dense = np.arange(h_start, h_end, dz)

    ## 1) Compute virtual temperature for environment and parcel
    # environment
    w_interp = interp1d(sounding.height, sounding.mixing_ratio, 'slinear')
    T_interp = interp1d(sounding.height, sounding.temperature, 'slinear')
    w_env = w_interp(z_dense)
    T_env = T_interp(z_dense)
    Tv_env = T_env*(1. + 0.61*w_env*1e-3)

    # parcel
    z_to_p = interp1d(sounding.height, sounding.pressure, 'slinear')
    p_to_z = interp1d(sounding.pressure[::-1], sounding.height, 'slinear')
    z_mapped_p = z_to_p(z_dense)
    z_subs = z_dense[z_mapped_p > 110]
    z_mapped_p = z_mapped_p[z_mapped_p > 110]

    T_interp = interp1d(parcel_sounding.pressure[::-1], parcel_sounding.temperature[::-1], 'slinear')
    T_par = T_interp(z_mapped_p)

    P_lcl = parcel_sounding.misc["P_lcl"]
    Z_lcl = p_to_z(P_lcl)
    w_par = np.zeros_like(T_par)
    w_par = 0.622*calc_es(T_par)/(100.*z_mapped_p)
    w_par[z_mapped_p > P_lcl] = 0.

    Tv_par = T_par*(1. + 0.61*w_par)

    # Collect functions
    Tv_par_f = interp1d(z_subs, Tv_par, 'slinear')
    Tv_env_f = interp1d(z_dense, Tv_env, 'slinear')

    # Find zeros of the buoyancy term
    Tv_diff = Tv_par_f(z_subs) - Tv_env_f(z_subs)
    pairs = []
    left_sign = []
    for i in xrange(1, len(Tv_diff)):
        if Tv_diff[i]*Tv_diff[i-1] < 0:
            pairs.append((i-1, i))
            if Tv_diff[i-1] < 0: 
                left_sign.append("neg")
            else:
                left_sign.append("pos")

    # Assume that the last two pairs bracket the LFC and EL, respectively, and take
    # the middle point between each pairs' values to use as the z-coordinates of those
    # quantities. Then, integrate to calculate CAPE
    LFC_pair, EL_pair = pairs[-2:]
    LFC_z = 0.5*np.sum([z_subs[i] for i in LFC_pair])
    LFC_p = float(z_to_p(LFC_z))
    EL_z = 0.5*np.sum([z_subs[i] for i in EL_pair])
    EL_p = float(z_to_p(EL_z))

    z_cape_int = z_subs[(z_subs < EL_z) & (z_subs > LFC_z)]
    CAPE = g*trapz((Tv_par_f(z_cape_int)+273.15)/(Tv_env_f(z_cape_int)+273.15) - 1., 
                   z_cape_int)

    parcel_sounding.misc["CAPV"] = CAPE
    parcel_sounding.misc["EL_z"] = EL_z
    parcel_sounding.misc["EL_p"] = EL_p
    parcel_sounding.misc["LFC_z"] = LFC_z
    parcel_sounding.misc["LFC_p"] = LFC_p

    return parcel_sounding

if __name__ == "__main__":

    import os

    plt.close('all')

    ## Test the plotting routines
    data_dir = "/Users/gutmann/Desktop"
    filename = os.path.join(data_dir, "dnr_sounding.txt")
    #filename = os.path.join(data_dir, "kshv_20110425_18z.snd")
    sounding = read_sounding(filename)

    ## Skew-T
    fig = plt.figure(figsize=(8,8))
    ax = fig.add_subplot(111)
    ax, ps = plot_skewt(sounding, ax=ax, title="FWD - Ft Worth\nApril 25, 2011 - 18Z ",
                       lift_parcel=True, plot_winds=True, diags=["CAPE", ],
                       label_altitudes=True)
    # parcel=compute_CAPE(sounding)
    # print(parcel.misc["CAPV"])
    ## Hodograph
    # ax = plot_hodograph(sounding)
    plt.show()
    ex = raw_input("\ncontinue...")
