#!/usr/bin/env python

"""
A CGI  script to compute surface temperature that gives
top-of-atmosphere radiative balance, using CliMT.
Rodrigo Caballero, Mike Steder, March 2004
"""

#### SET THESE:
# where CliMT lives:
CliMTDir = '/home/rca/lib/python2.3/site-packages'

# where matplotlib lives:
MatplotlibDir = '/home/rca/lib/python2.3/site-packages'

# where matplotlib keeps its fonts etc:
Matplotlibdata = '/home/rca/share/matplotlib'

# where ouput images will be served from: 
ImageDir = '/var/www/cgidata'

# where to retrieve the image over http
ImageURL = 'http://geoflop.uchicago.edu/cgidata' 
#############

## Development Debugging:
import cgitb; cgitb.enable()
# For production:  Send output to a file in /tmp
# so you can debug errors without spewing error messages
# at your users.
# import cgitb; cgitb.enable(display=0, logdir="/tmp")

import cgi,sys,os,Numeric
# Put CliMT in search path and import
sys.path.append(CliMTDir)
import climt
# set MATPLOTLIB environment variable and import matplotlib
os.environ['MATPLOTLIBDATA'] = Matplotlibdata
os.environ['TTFPATH'] = Matplotlibdata
from matplotlib.pylab import *
import matplotlib

# set up time stamp
import time
TimeStamp = ''
for i in time.localtime():
  TimeStamp += str(i)

## Step 1:  Setup HTML Header (Content-Type)
print "Content-Type: text/html"     # Just set the standard html content type.
print                               # Blank line signifies the end of the header info.
print ""
print ""
print ""
print "
Results
"


## Step 2:  Get the form data:
form = cgi.FieldStorage()

## Step 3:  Check for errors or missing fields.
## Ex. Missing Data:
FIELDS = ["scheme","Insolation","Imbalance","Trop_height","lapse_rate","CO2",
          "Albedo","RH","RH_control","T_air_RH","Trop_height_RH",
          "Cloud_frac_lo","Cloud_water_lo","Drop_size",
          "Cloud_frac_hi","Cloud_water_hi"] 
QUIT = 0
empty_fields = []
for field in FIELDS:
  if not form.has_key(field):
    empty_fields.append(field)
    QUIT = 1
## This is where I would check to make sure the values are valid
    
## Generate some nice formatting:  "Error: Please fill in field1, field2, field3 and field4"
if len(empty_fields) == 1:
  print '
Error: Please fill in the ', empty_fields[0], 'field.'
elif len(empty_fields) > 1:
  print '
Error: Please fill in the ',empty_fields[0]
  del empty_fields[0]
  while len(empty_fields) > 1:
    print ',',empty_fields[0]
    del empty_fields[0]
  print 'and', empty_fields[0], ' fields.' 
## Ex. Dealing with lists, multiple fields with the same name:
#value = form.getvalue("username", "")
#if isinstance(value, list):
#    # Multiple username fields specified
#    usernames = ",".join(value)
#else:
#    # Single or no username field specified
#    usernames = value

## Ex. Uploading Files in Forms:
#fileitem = form["userfile"]
#if fileitem.file:
#    # It's an uploaded file; count lines
#    linecount = 0
#    while 1:
#        line = fileitem.file.readline()
#        if not line: break
#        linecount = linecount + 1

if QUIT == 1:
  print "
",
  print "Error:  An error occured processing your request."
  print "
"
  print ""
  print ""
  sys.exit(1)
  
## Step 4:  Access and process form data.
scheme         = str(form.getvalue("scheme"))
Insolation     = float(form.getvalue("Insolation"))
Imbalance      = float(form.getvalue("Imbalance"))
Albedo         = float(form.getvalue("Albedo"))/100.
Trop_height    = float(form.getvalue("Trop_height"))
lapse_rate     = float(form.getvalue("lapse_rate"))
CO2            = float(form.getvalue("CO2")) + 1.e-9
CH4            = float(form.getvalue("CH4"))+ 1.e-9
N2O            = float(form.getvalue("N2O"))+ 1.e-9
RH             = float(form.getvalue("RH")) / 100.
RH_control     = int(form.getvalue("RH_control"))
T_air_RH       = float(form.getvalue("T_air_RH")) + 273.15
Trop_height_RH = float(form.getvalue("Trop_height_RH"))
Cloud_frac_lo  = float(form.getvalue("Cloud_frac_lo"))/100.
Cloud_water_lo = float(form.getvalue("Cloud_water_lo"))
Drop_size      = float(form.getvalue("Drop_size"))
Cloud_frac_hi  = float(form.getvalue("Cloud_frac_hi"))/100.
Cloud_water_hi = float(form.getvalue("Cloud_water_hi"))
Dry_strat      = form.getvalue('Dry_strat')
if Dry_strat is None:
  Dry_strat=0
else:
  Dry_strat=1

## Step 5: use CliMT to do stuff
# instantiate radiation object, get number of levels
r=climt.radiation(scheme=scheme)
levs=r.nlev

# define some fixed profiles
ps = 1000.
p  = ( Numeric.arrayrange(levs, typecode='d')+ 0.5 ) * ps/levs
cloud_lev_hi = int(levs*0.2)  # put high cloud roughly at 200 mb
cloud_lev_lo = int(levs*0.8)  # put low cloud roughly at 800 mb
cldf = Numeric.zeros( levs, typecode='d')  # Cloud frac
clwp = Numeric.zeros( levs, typecode='d')  # Cloud liquid water path
cldf[cloud_lev_lo] = Cloud_frac_lo
cldf[cloud_lev_hi] = Cloud_frac_hi
clwp[cloud_lev_lo] = Cloud_water_lo
clwp[cloud_lev_hi] = Cloud_water_hi


input={}
input['solin'] = Insolation
input['r_liq'] = Drop_size
input['r_ice'] = Drop_size
input['aldir'] = Albedo
input['aldif'] = Albedo
input['asdir'] = Albedo
input['asdif'] = Albedo
input['co2'] = CO2
input['ch4'] = CH4
input['n2o'] = N2O
input['ps'] = ps
input['p'] = p
input['cldf'] = cldf
input['clwp'] = clwp

# define some functions
def profiles(Ts):
  # assume near-surface temp is 1 K less than surface
  T_air = Ts - 1.
  # scale height (assume dry atmos) see Holton pg. 21
  Tmean =  (T_air**2*Trop_height + lapse_rate*T_air*Trop_height**2
            + lapse_rate**2*Trop_height**3)/(T_air*Trop_height + lapse_rate*Trop_height**2/2)
  H = r.Params['Rd']*Tmean/r.Params['g'] * 1.e-3 # [km]

  Tmean_RH =  (T_air_RH**2*Trop_height + lapse_rate*T_air_RH*Trop_height**2
          + lapse_rate**2*Trop_height**3)/(T_air_RH*Trop_height + lapse_rate*Trop_height**2/2)
  H_RH = r.Params['Rd']*Tmean_RH/r.Params['g'] * 1.e-3 #  scale height [km]
  # now compute profiles
  z = -H*Numeric.log(p/1000.)
  T = T_air + lapse_rate*z
  q  = Numeric.zeros(levs, typecode='d')
  for k in range(levs-1,-1,-1): # compute from bottom up
      if z[k] < Trop_height:
          q[k] = climt.thermodyn.qs(T[k],p[k])*RH
      else:
          T[k] = T[k+1]
          if Dry_strat:
              q[k] = 1.e-9
          else:
              q[k] = 1.e-4 #climt.thermodyn.qsflatau(T[k],p[k],2)*RH
  if RH_control:
      z_RH= -H_RH * Numeric.log(p/1000.)
      T_RH = T_air_RH + lapse_rate*z_RH
      for k in range(levs-1,-1,-1): # compute from bottom up
          if (z_RH[k] < Trop_height_RH):            
              q[k] = climt.thermodyn.qs(T_RH[k],p[k])*RH
          else:
              T_RH[k] = T_RH[k+1]
              if Dry_strat:
                  q[k] = 1.e-9
              else:
                  q[k] = 1.e-4 #climt.thermodyn.qsflatau(T_RH[k],p[k],2)*RH
  return T, q, z

def TOAFlux(Ts):
    T,q,z = profiles(Ts)
    input['T'] = T 
    input['Ts'] = Ts 
    input['q'] = q 
    r(**input)
    return r['swflx'][0] + r['lwflx'][0] + Imbalance

# compute equilibrium surface temperature
Ts = 300. # initial guess
try:
 Ts = climt.mathutil.ridder_root(TOAFlux, (173.15,353.15), accuracy=0.5)
except climt.mathutil.BracketingException, err:
  if str(err) == 'initial interval does not bracket a root: root probably to the right of interval':
    print '
Runaway greenhouse!'
  if str(err) == 'initial interval does not bracket a root: root probably to the left of interval':
    print '
Runaway icehouse!'
  sys.exit(1)
  
print
print 'Equilibrium near-surface air temperature is %4.1f oC (%4.1f K)' % ((Ts-273.15-1.),Ts-1.)
print

# compute profiles
T,q,z = profiles(Ts)
input['T'] = T 
input['Ts'] = Ts 
input['q'] = q 
r(**input)
swhr=r['swhr'] #*86400.
lwhr=r['lwhr'] #*86400.

## Step 6: make plot

from matplotlib.ticker import *
from matplotlib import rcParams
rcParams['tick.labelsize']  = 8

subplot(231)
plot(T-273.15,z, 'b-',linewidth=1)
#title(r'$\rm{Temperature} (^__\rm{o}\rm{C})$')
title('Temperature (C)',fontsize=10)
ylabel('height (km)',fontsize=10)
set(gca(), 'xlim', [-100,50],'ylim', [0,25])

subplot(232)
plot(q,z, 'b-',linewidth=1)
title('Specific humidity (g/kg)',fontsize=10)
set(gca(), 'xlim', [0,30],'ylim', [0,25])

subplot(233)
plot(cldf*100,z,'bo',clwp,z, 'ro',linewidth=1,markersize=2)
title('Cloud',fontsize=10)
legend(('fraction (%)', 'water path (g/m2)'), 'upper right')
set(gca(), 'xlim', [0,100],'ylim', [0,25])

ax=subplot(234)
ax.xaxis.set_major_locator(MultipleLocator(100))
plot(r['lwflx'][:-1],z, 'b-',linewidth=1)
title('Longwave flux (W/m2)',fontsize=10)
ylabel('height (km)',fontsize=10)
set(gca(), 'xlim', [-400,0],'ylim', [0,25])

ax=subplot(235)
ax.xaxis.set_major_locator(MultipleLocator(100))
plot(r.swflx[:-1],z, 'b-',linewidth=1)
title('Shortwave flux (W/m2)',fontsize=10)
set(gca(), 'xlim', [0,400],'ylim', [0,25])

subplot(236)
plot(lwhr,z,'b-', swhr,z,'r-', swhr+lwhr,z,'k-',linewidth=1)
title('Heating rates (K/day)',fontsize=10)
legend(('LW', 'SW', 'Total'), 'upper right')
set(gca(), 'xlim', [-4,4],'ylim', [0,25])

savefig(os.path.join(ImageDir,TimeStamp),dpi=100)
 
## Step 7: send stuff back to client
print ''

print '
'
print 'Profiles'
print("lev    p      z      T       q       LW    LW heating  SW  SW heating cld frac cld water\n")
for i in range(r.nlev):
    print("%3i %6.1f %7.2f %6.1f %6.2f %10.2f %6.2f %10.2f %6.2f %6.1f %6.1f" % \
          (i, p[i], z[i], T[i]-273.15, q[i], r['lwflx'][i], lwhr[i], r['swflx'][i], swhr[i] , cldf[i]*100., clwp[i])) 
print '
'