#!/usr/bin/env python
from __future__ import print_function
# Run this script with pnc options
import os

from PseudoNetCDF.pncload import PNCConsole
from PseudoNetCDF.coordutil import getmap, getlatbnds, getlonbnds
from PseudoNetCDF.coordutil import getybnds, getxbnds
from PseudoNetCDF.sci_var import getvarpnc, slice_dim

import warnings
import numpy as np
import matplotlib.pyplot as plt
warn = warnings.warn
pl = plt
LogFormatter = plt.matplotlib.ticker.LogFormatter
BoundaryNorm = plt.matplotlib.colors.BoundaryNorm
# from PseudoNetCDF.plotutil import *


def makemaps(args):
    ifiles = args.ifiles
    cbar = None
    ifile = ifiles[0]
    if args.iter != []:
        ifile, = ifiles
        ifiles = []
        for dimk in args.iter:
            ifiles += [slice_dim(getvarpnc(ifile, None), '%s,%d' % (dimk, i))
                       for i in range(len(ifile.dimensions[dimk]))]
    ax = plt.gca()
    map = getmap(ifile, resolution=args.resolution)
    if args.coastlines:
        map.drawcoastlines(ax=ax)
    if args.countries:
        map.drawcountries(ax=ax)
    if args.states:
        map.drawstates(ax=ax)
    if args.counties:
        map.drawcounties(ax=ax)
    for si, shapefile in enumerate(args.shapefiles):
        shapeopts = shapefile.split(',')
        shapepath = shapeopts[0]
        shapeoptdict = eval('dict(' + ','.join(shapeopts[1:]) + ')')
        shapename = os.path.basename(shapepath)[:-3] + str(si)
        map.readshapefile(shapepath, shapename, ax=ax, **shapeoptdict)
    args.map = map
    fig = plt.gcf()
    if len(args.figure_keywords) > 0:
        plt.setp(fig, **args.figure_keywords)

    ax = plt.gca()
    if len(args.axes_keywords) > 0:
        plt.setp(ax, **args.axes_keywords)

    map = args.map
    nborders = len(ax.collections)
    for fi, ifile in enumerate(ifiles):
        if map.projection in ('lcc', 'merc'):
            lat = ifile.variables['latitude']
            lon = ifile.variables['longitude']
            latb, latunit = getybnds(ifile)[:]
            lonb, lonunit = getxbnds(ifile)[:]
        else:
            lat = ifile.variables['latitude']
            lon = ifile.variables['longitude']
            latb, latunit = getlatbnds(ifile)[:]
            lonb, lonunit = getlonbnds(ifile)[:]

        if latb.ndim == lonb.ndim and lonb.ndim == 2:
            LON, LAT = lonb, latb
        else:
            LON, LAT = np.meshgrid(lonb.view(np.ndarray),
                                   latb.view(np.ndarray))

        variables = args.variables
        if variables is None:
            def isgeo(var):
                geo2d = set(['latitude', 'longitude'])
                vard = getattr(var, 'coordinates', '').split()
                hasgeo2d = len(geo2d.intersection(vard)) == 2
                return hasgeo2d
            variables = [key for key, var in ifile.variables.items()
                         if isgeo(var)]
        if len(variables) == 0:
            raise ValueError('Unable to heuristically determin plottable ' +
                             'variables; use -v to specify variables for ' +
                             'plotting')
        for varkey in variables:
            ax = plt.gca()

            if not args.overlay:
                del ax.collections[nborders:]
            var = ifile.variables[varkey]
            if args.squeeze:
                vals = var[:].squeeze()
            else:
                vals = var[:]
            vmin, vmax = vals.min(), vals.max()
            if args.normalize is None:
                from scipy.stats import normaltest
                if normaltest(vals.ravel())[1] < 0.001:
                    cvals = np.ma.compressed(vals)
                    boundaries = np.percentile(cvals, np.arange(0, 110, 10))
                    warn('Autoselect deciles colormap of %s; override ' +
                         'width --norm' % varkey)
                else:
                    boundaries = np.linspace(vmin, vmax, num=11)
                    warn(('Autoselect linear colormap of %s; override ' +
                          'width --norm') % varkey)
                ordermag = (boundaries.max() /
                            np.ma.masked_values(boundaries, 0).min())
                if (ordermag) > 10000:
                    formatter = LogFormatter(labelOnlyBase=False)
                else:
                    formatter = None
                norm = BoundaryNorm(boundaries, ncolors=256)
            else:
                norm = eval(args.normalize)
                formatter = None
            if args.colorbarformatter is not None:
                try:
                    formatter = eval(args.colorbarformatter)
                except Exception:
                    formatter = args.colorbarformatter

            if norm.vmin is not None:
                vmin = norm.vmin
            if norm.vmax is not None:
                vmax = norm.vmax
            varunit = getattr(var, 'units', 'unknown').strip()
            if args.verbose > 0:
                print(varkey, sep='')
            if vals.ndim == 1:
                notmasked = ~(np.ma.getmaskarray(lon[:]) |
                              np.ma.getmaskarray(lat[:]) |
                              np.ma.getmaskarray(vals[:]))
                scatlon = lon[:][notmasked]
                scatlat = lat[:][notmasked]
                scatvals = vals[:][notmasked]
                patches = map.scatter(scatlon[:], scatlat[:], c=scatvals,
                                      edgecolors='none', s=24, norm=norm,
                                      ax=ax, zorder=2)
            else:
                if vals.ndim != 2:
                    dimlendictstr = str(dict(zip(var.dimensions, var.shape)))
                    warn('Maps require 2-d data; values right now %s %s' %
                         (str(vals.shape), dimlendictstr))
                patches = map.pcolor(LON, LAT, vals, norm=norm, ax=ax)
            if lonunit == 'x (m)':
                ax.xaxis.get_major_formatter().set_scientific(True)
                ax.xaxis.get_major_formatter().set_powerlimits((-3, 3))
            if latunit == 'y (m)':
                ax.yaxis.get_major_formatter().set_scientific(True)
                ax.yaxis.get_major_formatter().set_powerlimits((-3, 3))
            ax.set_xlabel(lonunit)
            ax.set_ylabel(latunit)
            height = np.abs(np.diff(ax.get_ylim()))
            width = np.abs(np.diff(ax.get_xlim()))
            if width >= height:
                orientation = 'horizontal'
            else:
                orientation = 'vertical'
            if cbar is None:
                cax = None
            else:
                cax = cbar.ax
                cax.cla()

            if vals.max() > vmax and vals.min() < vmin:
                extend = 'both'
            elif vals.max() > vmax:
                extend = 'max'
            elif vals.min() < vmin:
                extend = 'min'
            else:
                extend = 'neither'
            cbar = plt.gcf().colorbar(patches, orientation=orientation,
                                      cax=cax, extend=extend, format=formatter,
                                      spacing='proportional')
            del cbar.ax.texts[:]
            varminmaxtxt = ('; min=%.3g; max=%.3g)' %
                            (var[:].min(), var[:].max()))
            cbar.set_label(varkey + ' (' + varunit + varminmaxtxt)
            # if orientation == 'vertical':
            #     cbar.ax.text(.5, 1.05, '%.3g' % var[:].max(),
            #                  horizontalalignment = 'center',
            #                  verticalalignment = 'bottom')
            #     cbar.ax.text(.5, -.06, '%.3g ' % var[:].min(),
            #                  horizontalalignment = 'center',
            #                  verticalalignment = 'top')
            # else:
            #     cbar.ax.text(1.05, .5, ' %.3g' % var[:].max(),
            #                  verticalalignment = 'center',
            #                  horizontalalignment = 'left')
            #     cbar.ax.text(-.06, .5, '%.3g ' % var[:].min(),
            #                  verticalalignment = 'center',
            #                  horizontalalignment = 'right')
            cbar.update_ticks()
            fmt = args.figformat
            outpath = args.outpath
            if len(ifiles) > 1:
                lstr = str(fi).rjust(len(str(len(ifiles))), '0')
                if args.verbose > 0:
                    print('adding numeric suffix for file', lstr)
            else:
                lstr = ''

            figpath = os.path.join(outpath + varkey + lstr + '.' + fmt)
            if args.interactive:
                csl = PNCConsole(locals=globals())
                csl.interact()
            for cmd in args.plotcommands:
                exec(cmd)
            plt.savefig(figpath)
            if args.verbose > 0:
                print('Saved fig', figpath)
