#! /usr/bin/env python

# Copyright (C) 2010
# Glimmer-CISM contributors - see AUTHORS file for list of contributors
#
# This file is part of Glimmer-CISM.
#
# Glimmer-CISM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or (at
# your option) any later version.
#
# Glimmer-CISM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Glimmer-CISM.  If not, see <http://www.gnu.org/licenses/>.
#
# Glimmer-CISM is hosted on BerliOS.de:
# https://developer.berlios.de/projects/glimmer-cism/

"""A simple plot of CF fields."""

import string
import PyGC
import matplotlib.pyplot
import matplotlib.cm, matplotlib.colors
from mpl_toolkits.axes_grid.anchored_artists import AnchoredText

have_basemap=True
try:
    from mpl_toolkits.basemap import Basemap
except:
    have_basemap=False
import numpy,math

def makeMask(mask,rgba):
    a = numpy.zeros(list(mask.shape)+[4])
    for i in range(0,3):
        a[:,:,i] = rgba[i]
    a[:,:,3] = numpy.where(mask,rgba[3],0)
    return a

# creating option parser
parser = PyGC.GCOptParser()
parser.variable()
parser.profile_file(plist=True)
parser.time()
parser.region()
parser.plot()
opts = PyGC.GCOptions(parser,1)

if opts.options.level == None:
    level = 0
else:
    level = opts.options.level

infile = opts.gcfile()
var  = opts.vars(infile)
time = opts.times(infile)
clip = opts.options.clip

plot_geo_coord = infile.projection!='lin' and have_basemap and opts.options.geo_coord

e = numpy.array([infile.ll_xy[0],infile.ur_xy[0],infile.ll_xy[1],infile.ur_xy[1]])
extent=numpy.array([var.xdim[0],var.xdim[-1],var.ydim[0],var.ydim[-1]])
if not plot_geo_coord:
    e = e/1000.
    extent=extent/1000.

fig = matplotlib.pyplot.figure(figsize=opts.plotsize)
m = fig.add_axes([0.1,0.1,0.8,0.8])

# indicate area above sea level
if opts.options.land:
    v = infile.getvar('topg').get2Dfield(time)
    l = numpy.zeros([v.shape[0],v.shape[1],4])+100
    l[:,:,3] = numpy.where(v>0,1,0)
    m.imshow(l,origin='lower',extent=extent)

# plot the data
norm,cmap,title = opts.colourmap(infile)
im = m.imshow(var.get2Dfield(time,level=level,clip=opts.options.clip),norm=norm,cmap=cmap,extent=extent,origin='lower')

# shade plot
if opts.options.illuminate!=None:
    ls = matplotlib.colors.LightSource()
    v = ls.shade(infile.getvar(opts.options.illuminate).get2Dfield(time),cmap=matplotlib.cm.gray)
    v[:,:,3] = (1-v[:,:,0])*0.8
    v[:,:,0:3] = 0.05
    m.imshow(v,origin='lower',extent=extent)

# plot velo vectors
if opts.options.vectors != None:
    n={'vel':'vel','bvel':'bas'}[opts.options.vectors]
    u = infile.getvar('u%s'%n).get2Dfield(time,level=level,clip=opts.options.vectors)
    v = infile.getvar('v%s'%n).get2Dfield(time,level=level,clip=opts.options.vectors)
    x,y = numpy.meshgrid(infile.file.variables['x0'],infile.file.variables['y0'])
    m.quiver(x,y,u,v,units='x')

m.set_aspect('equal')

if opts.options.legend!=None:
    if opts.options.legend=='v':
        orientation='vertical'
    else:
        orientation='horizontal'
    cb = fig.colorbar(im, extend='both', spacing='uniform',orientation=orientation, shrink=0.8)
    cb.set_label(title)

# plot profiles
if opts.options.profname !=None:
    i=0
    for pn in opts.options.profname:
        pdata = PyGC.IO.GCprofile(infile,interval=opts.options.interval)
        pdata.coords_file(pn,opts.options.prof_is_projected)
        if pdata.interpolated.shape[1] > 0:
            m.plot(pdata.interpolated[0],pdata.interpolated[1],color='k')
            m.annotate(string.ascii_uppercase[i],xy=pdata.interpolated[:,0])
        i=i+1


# plot a geographic coordinate system and coastlines
if plot_geo_coord:
    longs = numpy.arange(math.ceil(infile.minmax_long[0]/10.), math.floor(infile.minmax_long[1]/10.))*10.
    lats = numpy.arange(math.ceil(infile.minmax_lat[0]/10.), math.floor(infile.minmax_lat[1]/10.))*10.

    # turn projection parameters into parameters for setting up basemap plot
    params = {}
    for p in infile.projection.proj4_params():
        p = p.split('=')
        k=p[0]
        v=p[1]
        if k=='proj':
            params['projection'] = v
        elif k in ['ellps','y_0','x_0']:
            pass
        else:
            params[k] = float(v)
    params['llcrnrlon']=infile.ll_geo[0]
    params['llcrnrlat']=infile.ll_geo[1]
    params['urcrnrlon']=infile.ur_geo[0]
    params['urcrnrlat']=infile.ur_geo[1]
    params['resolution']='l'
    params['ax'] = m

    b = Basemap(**params)
    b.drawcoastlines()
    b.drawparallels(lats,labels=[1,1,0,1])
    b.drawmeridians(longs,labels=[1,1,0,1])

# plot mask
if opts.options.mask:
    mask = numpy.array(infile.getvar('thkmask').get2Dfield(time),int)
    adjustedMask = numpy.where(mask<0,0,mask)

    thin = makeMask(adjustedMask&1==1,[1.,1.,1.,0.7])
    m.imshow(thin,origin='lower',extent=extent)

    shelf = makeMask(adjustedMask&4==4,[0,0,1,0.8])
    m.imshow(shelf,origin='lower',extent=extent)

    interior = makeMask(adjustedMask&2==2,[1,1,1,0.8])
    m.imshow(interior,origin='lower',extent=extent)

    grounding = makeMask(adjustedMask&16,[0,0,0,1])
    m.imshow(grounding,origin='lower',extent=extent)
    
# plot time stamp
if opts.options.timestamp:
    m.add_artist(AnchoredText('%.2fka'%infile.time(time),loc=4))

# display title
if opts.options.title:
    matplotlib.pyplot.title(infile.title)

m.set_xlim(e[:2])
m.set_ylim(e[2:])

if opts.options.output==None:
    matplotlib.pyplot.show()
else:
    matplotlib.pyplot.savefig(opts.options.output)
