#!/usr/bin/env python

try:
    from netCDF4 import Dataset
except:
    print "netCDF4 is not installed!"
    sys.exit(1)

import numpy as np
import pylab as plt

from optparse import OptionParser

parser = OptionParser()
parser.usage = "usage: %prog [options] FILE"
parser.description = "A script to compare PISM flowline velocities with Stokes solution."

(options, args) = parser.parse_args()

plot_acab = True

if len(args) != 1:
    print('wrong number of arguments, 1 expected')
    exit(1)

try:
    nc = Dataset(args[0], 'r')
except:
    print(("file %s not found ... ending ..." % args[0]))
    exit(2)


def permute(variable, output_order=('time', 'z', 'zb', 'y', 'x')):
    """Permute dimensions of a NetCDF variable to match the output storage order."""
    input_dimensions = variable.dimensions

    # filter out irrelevant dimensions
    dimensions = filter(lambda(x): x in input_dimensions,
                        output_order)

    # create the mapping
    mapping = map(lambda(x): dimensions.index(x),
                  input_dimensions)

    if mapping:
        return np.transpose(variable[:], mapping)
    else:
        return variable[:]              # so that it does not break processing "mapping"


x = nc.variables["x"][:]
b = np.squeeze(nc.variables["topg"][:])
s = np.squeeze(nc.variables["usurf"][:])
h = np.squeeze(nc.variables["thk"][:])
z = nc.variables["z"][:]

mask = np.zeros_like(h)
mask[h <= 1] = 1

us = np.ma.array(data=np.squeeze(nc.variables["uvelsurf"][:]), mask=mask)
ub = np.ma.array(data=np.squeeze(nc.variables["uvelbase"][:]), mask=mask)
# stuff needed for contour plots
xx = (np.tile(x, [len(z), 1]))
zz = ((np.tile(z, [len(x), 1])).transpose() + b)

# ignore the first level
cts = np.squeeze(permute(nc.variables["cts"]))
liqfrac = np.squeeze(permute(nc.variables["liqfrac"]))
temppa = np.squeeze(permute(nc.variables["temp_pa"]))
mask2 = np.zeros_like(cts)
mask2[zz > s] = 1
cts = np.ma.array(data=cts, mask=mask2)
liqfrac = np.ma.array(data=liqfrac, mask=mask2)
temppa = np.ma.array(data=temppa, mask=mask2)

# Contour level of the CTS
cts_level = [1, 1]
liqfrac_levels = np.arange(0, 2.5, .25)
temppa_levels = [-6, -5, -4, -3, -2, -1, -.0001]

fig = plt.figure(figsize=(6.4, 7.4))

axUpperLeft = plt.axes([0.1, 0.6, 0.8, 0.25])
axLower = plt.axes([0.1, 0.05, 0.8, 0.5])

axUpperLeft.plot(x, us, color='#377EB8', lw=1.5)
axUpperLeft.plot(x, ub, '--', color='#377EB8', lw=1.5)
axUpperLeft.axes.set_xlim(-250, 3500)
axUpperLeft.axes.set_ylabel("velocity [m a$^{-1}$]")
plt.setp(axUpperLeft, xticks=[])
if (plot_acab == True):
    acab = np.squeeze(nc.variables["climatic_mass_balance"][:])
    axUpperRight = axUpperLeft.twinx()
    axUpperRight.plot(x, acab / 910.0, color='#984EA3', lw=1.5)
    axUpperRight.axes.set_ylabel("mass balance [m a$^{-1}$]")
    axUpperRight.axes.set_xlim(-250, 3500)
axLower.plot(x, b, color='black', lw=1.5)
axLower.plot(x, s, color='black', lw=1.5)
c1 = axLower.contourf(xx, zz, liqfrac * 100, liqfrac_levels, cmap=plt.cm.Reds)
plt.colorbar(mappable=c1, ax=axLower, orientation='horizontal', pad=0.05, shrink=0.75, extend="max")
c2 = axLower.contourf(xx, zz, temppa, temppa_levels, cmap=plt.cm.Blues_r, lw=1)
plt.colorbar(mappable=c2, ax=axLower, orientation='horizontal', ticks=[-6, -5, -4, -3, -2, -1, 0], pad=0.20, shrink=0.75)
axLower.contour(xx, zz, cts, cts_level, colors='black', linestyles='dashed', lw=1)
axLower.axes.set_xlim(-250, 3500)
axLower.axes.set_ylim(1100, 1800)
axLower.axes.set_xlabel("distance from bergschrund [m]")
axLower.axes.set_ylabel("elevation [m a.s.l.]")


plt.savefig('sg_results.pdf', bbox_inches='tight', pad_inches=0.35)

nc.close()
