import math
import numpy as np
import pylab as pl


# most of the code adapted from pywrfplot
# http://code.google.com/p/pywrfplot/

# purpose: plot raw data of Temperature and Dewpoint in skewT-log(p)-diagram

skewness = 37.5
# Defines the ranges of the plot, do not confuse with P_bot and P_top
P_b = 105000.
P_t = 10000. 
dp = 100.
plevs = np.arange(P_b,P_t-1,-dp)

T_zero = 273.15

# P_top must be the same as what is used in the WRF simulation
P_top = 10**4
P_bot = 10**5


L = 2.501e6 # latent heat of vaporization
R = 287.04  # gas constant air
Rv = 461.5  # gas constant vapor
eps = R/Rv

cp = 1005.
cv = 718.
kappa = (cp-cv)/cp
g = 9.81


# constants used to calculate moist adiabatic lapse rate
# See formula 3.16 in Rogers&Yau
a = 2./7.
b = eps*L*L/(R*cp)
c = a*L/R





def SkewTPlot():
    
    pl.figure()
    _isotherms()
    _isobars()
    _dry_adiabats()
    _moist_adiabats()
       
    _plot_data()
    
    
    pl.axis([-40,50,P_b,P_t])
    pl.xlabel('Temperatur ($^{\circ}\! C$)')
    xticks = np.arange(-40,51,5)
    pl.xticks(xticks,['' if tick%10!=0 else str(tick) for tick in xticks])
    pl.ylabel('Pressure (hPa)')
    yticks = np.arange(P_bot,P_t-1,-10**4)
    pl.yticks(yticks,yticks/100)
    
    pl.show()
    pl.close()
    pl.clf()


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


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

def _isobars():
    for n in np.arange(P_bot,P_t-1,-10**4):
        pl.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)
        pl.semilogy(dry_adiabat, plevs, basey=math.e, color = 'brown', \
                     linestyle='dashed', linewidth = .5)



def es(T):
    """Returns saturation vapor pressure (Pascal) at temperature T (Celsius)
    Formula 2.17 in Rogers&Yau"""
    return 611.2*np.exp(17.67*T/(T+243.5))


def gamma_s(T,p):
    """Calculates moist adiabatic lapse rate for T (Celsius) and p (Pa)
    Note: We calculate dT/dp, not dT/dz
    See formula 3.16 in Rogers&Yau for dT/dz, but this must be combined with
    the dry adiabatic lapse rate (gamma = g/cp) and the 
    inverse of the hydrostatic equation (dz/dp = -RT/pg)"""
    esat = es(T)
    wsat = eps*esat/(p-esat) # Rogers&Yau 2.18
    numer = a*(T+T_zero) + c*wsat
    denom = p * (1 + b*wsat/((T+T_zero)**2)) 
    return numer/denom # Rogers&Yau 3.16



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))
        pl.semilogy(moist_adiabat, ps, basey=math.e, color = 'green', \
                     linestyle = 'dashed', linewidth = .5)



def _plot_data():
    p,h,T,Td = np.loadtxt('/Users/daniel/tmp/soundingdata.txt',usecols=range(0,4),unpack=True)
    
    pl.semilogy(T+ _skewnessTerm(p*100),p*100,basey=math.e, color=('black'),linestyle=('solid'),linewidth= 1.5)
    
    pl.semilogy(Td+ _skewnessTerm(p*100),p*100,basey=math.e, color=('black'),linestyle=('solid'),linewidth= 1.5)
    
    return 1


SkewTPlot()