#!/usr/bin/env Python

# waterLand.py
# Python script to display time series of wet/dry land as images.

# Author: Carl Drews
# Created: September 2008
# Copyright 2008 by Carl Drews

# This script was written for my Master's Thesis in Atmospheric
# and Oceanic Sciences, University of Colorado at Boulder.
# Use ocean_his.nc as the input file (from the Lake of Tanis case study).

import sys
from array import array
import os
import Nio
import matplotlib
from pylab import *
import numpy



cyanFraction = 0.85	# controls lightness of cyan;
			# use 1.0 for pure white

# Map depth of water to RGB or RGBA color.
# waterDepth in meters
# maxDepth deepest value anywhere in the domain
# alpha = also include the opacity
# return normalized RGB tuple to display water depth
def depthColor(waterDepth, maxDepth, alpha):
   # clip water depth to the maximum value
   if (waterDepth > maxDepth):
      waterDepth = maxDepth

   if (not alpha):
      rgb = numpy.zeros([3], 'f')
   else:
      rgb = numpy.zeros([4], 'f')

   rgb[0] = cyanFraction * (1.0 - waterDepth / maxDepth)
   rgb[1] = 1.0 - waterDepth / maxDepth
   rgb[2] = 1.0

   if (alpha):
      rgb[3] = 1.0

   if (False):  # black-and-white
      rgb[:] = 1.0      # white

   return rgb



# Map topography to RGB color.
# topoHeight in meters above mean sea level
# maxTopo highest elevation anywhere
# alpha = also include the opacity
# return normalized RGB tuple to display topography
def topoColor(topoHeight, maxTopo, alpha):
   # clip topography to the maximum value
   if (topoHeight > maxTopo):
      topoHeight = maxTopo

   if (not alpha):
      rgb = numpy.zeros([3], 'f')
   else:
      rgb = numpy.zeros([4], 'f')

   rgb[0] = min(2 * topoHeight / maxTopo, 1.0)
   rgb[1] = min(2 * (1 - topoHeight / maxTopo), 1.0)
   rgb[2] = 0.0

   if (False):  # black-and-white
      rgb[:] = 0.6      # medium gray

   return rgb



# Map bathymetry to RGB color.
# bathyHeight in meters below mean sea level (negative)
# maxBathy lowest exposed mud flats anywhere
# alpha = also include the opacity
# return normalized RGB tuple to display bathymetry
def bathyColor(bathyHeight, maxBathy, alpha):
   # clip bathymetry to the maximum value
   if (bathyHeight < maxBathy):
      bathyHeight = maxBathy

   if (not alpha):
      rgb = numpy.zeros([3], 'f')
   else:
      rgb = numpy.zeros([4], 'f')

   # start with burlywood (shallow) and move toward saddleBrown (deep)
   shallow = array([222.0, 184.0, 135.0]) / 255
   deep = array([139.0, 069.0, 019.0]) / 255

   rgb[0] = (bathyHeight / maxBathy) * (deep[0] - shallow[0]) + shallow[0]
   rgb[1] = (bathyHeight / maxBathy) * (deep[1] - shallow[1]) + shallow[1]
   rgb[2] = (bathyHeight / maxBathy) * (deep[2] - shallow[2]) + shallow[2]

   if (False):  # black-and-white
      rgb[:] = 0.8      # light gray

   return rgb



# Create section of KML file for Google Earth,
# animating transparent images over the Kedua Gap.
# You can pasted the fragment into a larger KML file.
# kmlName = name of the .kml file to create
def kmlFragment(kmlName):
   file = open(kmlName, 'w')

   # loop through all the time steps
   step = 10
   for t in range(0, 241, step):
      hour = (t * 6) / 60
      hourStr = "%(#)02d" % { "#": hour }
      nextHour = ((t+step) * 6) / 60
      nextHourStr = "%(#)02d" % { "#": nextHour }

      minute = (t * 6) % 60
      minuteStr = "%(#)02d" % { "#": minute }
      nextMinute = ((t+step) * 6) % 60
      nextMinuteStr = "%(#)02d" % { "#": nextMinute }

      file.write("<GroundOverlay>\n")
      file.write("  <name>Water-Land mask at " + hourStr + ":" + minuteStr + "</name>\n")
      file.write("  <description>Displays regions of water and land before the crossing.</description>\n")
      file.write("  <color>70ffffff</color>           <!-- opacity -->\n")
      file.write("  <TimeSpan><begin>2009-04-17T" + hourStr + ":" + minuteStr + ":00</begin>\n")
      file.write("    <end>2009-04-17T" + nextHourStr + ":" + nextMinuteStr + ":00</end></TimeSpan>\n")

      file.write("  <Icon>\n")
      file.write("    <href>KeduaGap%(#)03d.png</href>\n" % { "#": t})
      file.write("  </Icon>\n")
      file.write("  <LatLonBox>\n")
      file.write("    <north>31.2</north>\n")
      file.write("    <south>30.8</south>\n")
      file.write("    <east>32.6</east>\n")
      file.write("    <west>32.2</west>\n")
      file.write("  </LatLonBox>\n")
      file.write("</GroundOverlay>\n\n")

   file.close()

   return



# Main program begins here.

# Set useAlpha = True to include opacity as RGBA.
# This is useful for creating overlay images.
# Set alpha = False for standard RGB images.
useAlpha = False

# Process command-line arguments.
wantRomsFile = False
romsFile = "default.nc"

for arg in sys.argv:
   if (wantRomsFile):
      romsFile = arg
      wantRomsFile = False

   elif (arg == "-romsfile"):
      wantRomsFile = True

# let the log file know what's happening
print "ROMS file = ", romsFile

if (True):
   # test creation of color figures
   temp = numpy.zeros([2, 2, 3], 'f')
   temp[0,0,0] = 0.9
   temp[0,1,1] = 0.8
   temp[1,0,2] = 0.7
   temp[1,1,0] = 0.9
   temp[1,1,1] = 0.9
   temp[1,1,2] = 1.0
   imshow(temp, interpolation='nearest', origin='lower')
   colorbar()
   savefig("temp")
   close()
   clf()

# set up the bounds for the plot
# regional
lowerLeftY = 0
lowerLeftX = 0
upperRightY = 1201
upperRightX = 1801

# crossing detail
if (True):
   lowerLeftY = 3 * 120
   lowerLeftX = 7 * 120
   upperRightY = lowerLeftY + 4 * 120 + 1
   upperRightX = lowerLeftX + 4 * 120 + 1

# open the NetCDF file for reading
print "Nio version ", Nio.__version__
file = Nio.open_file(romsFile, 'r')

# extract the bathymetry and topography
# ROMS h refers to positive depth of ocean bottom.
heightName = "h"
height = file.variables[heightName][lowerLeftY:upperRightY, lowerLeftX:upperRightX] + 0.0

# calculate extreme values for color scaling
maxTopo = -min(height.flat)
print "maximum topography = ", maxTopo
maxBathy = -max(height.flat)
print "maximum bathymetry = ", maxBathy


# extract the free-surface elevation
zetaName = "zeta"
zeta = file.variables[zetaName][:,lowerLeftY:upperRightY,lowerLeftX:upperRightX] + 0.0

# extract the wet-dry mask
wetDryMask = file.variables["wetdry_mask_rho"][:,lowerLeftY:upperRightY,lowerLeftX:upperRightX] + 0.0

# extract the simulation time
oceanTime = file.variables["ocean_time"]

# calculate the maximum water depth for color scaling
allDepth = numpy.zeros(zeta.shape, 'f')
for t in range(0, zeta.shape[0]):
   allDepth[t] = (zeta[t] + height) * wetDryMask[t]
maxDepth = allDepth.max()
print "maximum water depth = ", maxDepth
sys.stdout.flush()

# set color scaling to show detail near sea level
maxTopo = 20.0
maxBathy = -10.0
maxDepth = 4.0

# create two color images the same size as wet_dry_mask
if (not useAlpha):
   baseImage = numpy.zeros([wetDryMask.shape[1], wetDryMask.shape[2], 3], 'f')
   colorImage = numpy.zeros([wetDryMask.shape[1], wetDryMask.shape[2], 3], 'f')
else:
   baseImage = numpy.zeros([wetDryMask.shape[1], wetDryMask.shape[2], 4], 'f')
   colorImage = numpy.zeros([wetDryMask.shape[1], wetDryMask.shape[2], 4], 'f')

# the topography and bathymetry never change, so set them up now
for j in range(0, baseImage.shape[0]):
   for i in range(0, baseImage.shape[1]):
      if (height[j,i] <= 0.0):
         # topography above sea level is green-yellow-red
         baseImage[j,i] = topoColor(-height[j,i], maxTopo, useAlpha)
      if (height[j,i] > 0.0):
         # bathymetry below sea level is brown-dark brown
         baseImage[j,i] = bathyColor(-height[j,i], maxBathy, useAlpha)

# create green-yellow-red color bar for land height
greenRedSegments = {
      'red':   [(0.0,  0.0, 0.0),
                (0.5,  1.0, 1.0),
                (1.0,  1.0, 1.0)],
      'green': [(0.0,  1.0, 1.0),
                (0.5,  1.0, 1.0),
                (1.0,  0.0, 0.0)],
      'blue':  [(0.0,  0.0, 0.0),
                (0.5,  0.0, 0.0),
                (1.0,  0.0, 0.0)]}

greenRedMap = matplotlib.colors.LinearSegmentedColormap(
   'white-blue', greenRedSegments)

# create cyan-blue color bar for water depth
cyanBlueSegments = {
      'red':   [(0.0,  cyanFraction, cyanFraction),
                (1.0,  0.0, 0.0)],
      'green': [(0.0,  1.0, 1.0),
                (1.0,  0.0, 0.0)],
      'blue':  [(0.0,  1.0, 1.0),
                (1.0,  1.0, 1.0)]}

cyanBlueMap = matplotlib.colors.LinearSegmentedColormap(
   'white-blue', cyanBlueSegments)

# create a time series of color images
for t in range(0, zeta.shape[0], 1):
   # display progress
   print t,
   sys.stdout.flush()
   colorImage[:] = baseImage[:]

   # set color according to water depth, topography, or bathymetry
   for j in range(0, colorImage.shape[0]):
      for i in range(0, colorImage.shape[1]):
         # Mask out the splatter less than 1 cm above DCRIT.
         if (wetDryMask[t,j,i] and allDepth[t,j,i] > 0.11):
            # calculate water depth
            colorImage[j,i] = depthColor(allDepth[t,j,i], maxDepth, useAlpha)

   # It seems kind of stupid to call imshow() twice here,
   # but this is what works.  Carl Drews - April 9, 2009

   # create the image with colorbar for water depth
   imshow(colorImage[:, :],
      interpolation='nearest', origin='lower',
      cmap=cyanBlueMap, vmin=0.0, vmax=maxDepth)
   if (True):
      colorbar(shrink=0.6, pad=-0.05)

   # redraw the image with colorbar for land height
   imshow(colorImage[:, :],
      interpolation='nearest', origin='lower',
      cmap=greenRedMap, vmin=0.0, vmax=maxTopo)
   if (True):
      colorbar(shrink=0.6)

   # label the axes
   ax = gca()		# get current axes

   # create minute label with leading zero
   minute = str(int(oceanTime[t] % 3600) / 60)
   if len(minute) == 1:
      minute = "0" + minute
   #minute = "00"

   if (True):
      ax.set_title("Lake of Tanis at "
         + str(int(oceanTime[t]/3600)) + ":" + minute
         + " hours")
      ax.set_xlabel("East Longitude")
      ax.set_ylabel("North Latitude")

   # For some reason the first tick label gets ignored.
   degree = u"\u00b0"

   # for regional view
   if (False):
      ax.set_xticks(range(0, 1800+1, 120))
      ax.set_xticklabels(["",
         "31.6" + degree, "",
         "31.8" + degree, "",
         "32.0" + degree, "",
         "32.2" + degree, "",
         "32.4" + degree, "",
         "32.6" + degree, "",
         "32.8" + degree, "",
         "33.0" + degree])

      ax.set_yticks(range(0, 1200+1, 120))
      ax.set_yticklabels(["",
         "30.6" + degree, "",
         "30.8" + degree, "",
         "31.0" + degree, "",
         "31.2" + degree, "",
         "31.4" + degree])

   # for crossing detail
   if (True):
      ax.set_xticks(range(0, 4*120+1, 120))
      ax.set_xticklabels([
         "32.2" + degree,
         "32.3" + degree,
         "32.4" + degree,
         "32.5" + degree,
         "32.6" + degree])

      ax.set_yticks(range(0, 4*120+1, 120))
      ax.set_yticklabels([
         "30.8" + degree,
         "30.9" + degree,
         "31.0" + degree,
         "31.1" + degree,
         "31.2" + degree])

   # save image to a file
   filename = "images/topo%(#)03d" % { "#": t }
   #savefig(filename, dpi=300)
   savefig(filename, transparent=useAlpha)
   close()
   clf()

# Close the input file.
file.close()

# create kml fragment for the ground overlays
kmlFragment("fragment.kml")

# return and report success
sys.exit(0)