#! /usr/bin/env python

# Copyright (C) 2005, 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/

"""Plot sediment maps"""

import PyGC,sys,os,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

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

if opts.options.colourmap != None:
    erosioncpt = opts.options.colourmap
else:
    erosioncpt = 'erosion.cpt'

infile = opts.gcfile()

time = opts.times(infile)
var = infile.getvar("seds3")

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
cmap = PyGC.GCcolourmap(infile.getvar("seds3"),filename=erosioncpt)
norm,cmap = cmap.norm,cmap.colourmap
title = 'change in sediment thickness [m]'

init_seds = infile.getvar("seds3").get2Dfield(0)
seds1 = infile.getvar("seds1").get2Dfield(time)
seds2 = infile.getvar("seds2").get2Dfield(time)
seds3 = infile.getvar("seds3").get2Dfield(time)
seds3 = seds3+seds1+seds2-init_seds

im = m.imshow(seds3,norm=norm,cmap=cmap,extent=extent,origin='lower')

# plot ice extent
thk = None
try:
    thk = infile.getvar("thk").get2Dfield(time)
except:
    pass
if thk!=None:
    m.contour(thk,[1],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)

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 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)

