#!/usr/bin/env python
"""

Pseudocolor plots of unstructured triangular grids.

Draws grid bathymetry or whatever is in column 4 of a grid file

"""
import os
import getopt
import math

from crane import plt
from crane import np
import crane.files.grid as gr

import matplotlib.tri as tri
from matplotlib.lines import Line2D
import cProfile


def signa(x1, x2, x3, y1, y2, y3):    # area formed by 3 points
    """Compute the signed area of an element given the 3 vertices"""
    return ((x1 - x3) * (y2 - y3) - (x2 - x3) * (y1 - y3)) * 0.5


def computeER(area):
    """Compute the equivalent radius given the area"""
    return 2.0 * math.sqrt(area) / np.pi


def computeAreas(tri):
    areas = tri.x * 0.0
    x1 = tri.x[tri.triangles[:, 0]]
    x2 = tri.x[tri.triangles[:, 1]]
    x3 = tri.x[tri.triangles[:, 2]]
    y1 = tri.y[tri.triangles[:, 0]]
    y2 = tri.y[tri.triangles[:, 1]]
    y3 = tri.y[tri.triangles[:, 2]]
    areas = signa(x1, x2, x3, y1, y2, y3)
    #areas = ((x1-x3)*(y2-y3)-(x2-x3)*(y1-y3)) * 0.5
    return areas


def mapscale(x, y, len, units):
    """ Draw a simple map scale """
    inv = plt.gca().transLimits.inverted()
    (xloc, yloc) = inv.transform((x, y))
    (xloclab, yloclab) = inv.transform((x, y - 0.005))
    l = Line2D([xloc, xloc + len], [yloc, yloc], lw=3., color='k')
    plt.text(xloclab, yloclab, '0', fontsize=10, ha='center', va='top')
    if units == 'km':
        dlen = ('''%.0f %s''') % (len / 1000.0, units)
    plt.text(xloclab + len, yloclab, dlen, fontsize=10, ha='center', va='top')
    plt.gca().add_line(l)


def plotGridScalarNodes(triang, z, cmap, vmin=None, vmax=None, label=None):
    """ Pseudocolor scalar at the nodes and add a scalebar"""
    plt.tripcolor(
        triang,
        z,
        shading='gouraud',
        cmap=cmap,
        vmin=vmin,
        vmax=vmax)
    cb = plt.colorbar(orientation='vertical', shrink=0.5)
    cb.set_clim(vmin, vmax)
    cb.ax.invert_yaxis()
    if label is not None:
        cb.set_label(label)
    for t in cb.ax.get_yticklabels():
        t.set_fontsize(10)


def plotGridScalarElements(triang, ze, cmap, vmin=None, vmax=None, label=None):
    """ Psuedocolor flat shaded elements"""
    plt.tripcolor(
        triang,
        facecolors=ze,
        shading='flat',
        cmap=cmap,
        vmin=vmin,
        vmax=vmax)
    cb = plt.colorbar(orientation='vertical', shrink=0.5)
    cb.set_clim(vmin, vmax)
    if label is not None:
        cb.set_label(label)
    for t in cb.ax.get_yticklabels():
        t.set_fontsize(10)


def gridToTriangulation(grid):
    """Return a Triangulation from a grid object"""

    tris = np.zeros((len(grid.elems), 3))
    tris[:, 0] = grid.elems[:, 1] - 1
    tris[:, 1] = grid.elems[:, 2] - 1
    tris[:, 2] = grid.elems[:, 3] - 1
    return tri.Triangulation(grid.nodes[:, 1], grid.nodes[:, 2], tris)


def noTicks(ax):
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xticks([])
    ax.set_yticks([])


def drawElements(
        triang,
        ze,
        title,
        pngfile,
        w,
        h,
        xmin,
        xmax,
        ymin,
        ymax,
        vmin,
        vmax,
        cmap,
        label,
        mapx,
        mapy,
        maplen,
        units):
    """ Pseudocolor flat shaded elements"""
    f = plt.figure(
        num=None,
        figsize=(
            w,
            h),
        dpi=100,
        facecolor='w',
        edgecolor='w')
    plt.gca().set_aspect('equal')
    plt.tripcolor(
        triang,
        facecolors=ze,
        shading='flat',
        cmap=cmap,
        vmin=vmin,
        vmax=vmax)
    cb = plt.colorbar(orientation='vertical', shrink=0.5)
    plt.gca().set_xlim(xmin, xmax)
    plt.gca().set_ylim(ymin, ymax)
    cb.set_clim(vmin, vmax)
    if label is not None:
        cb.set_label(label)
    for t in cb.ax.get_yticklabels():
        t.set_fontsize(10)
    noTicks(plt.gca())
    plt.title(title)
    f.tight_layout()
    mapscale(mapx, mapy, maplen, units)
    # plt.show()
    plt.savefig(pngfile, bbox_inches='tight')


def drawScalar(
        triang,
        z,
        title,
        pngfile,
        w,
        h,
        xmin,
        xmax,
        ymin,
        ymax,
        vmin,
        vmax,
        cmap,
        label,
        mapx,
        mapy,
        maplen,
        units):
    """ Pseudocolor scalar at the nodes and add a scalebar"""
    f = plt.figure(
        num=None,
        figsize=(
            w,
            h),
        dpi=100,
        facecolor='w',
        edgecolor='w')
    plt.gca().set_aspect('equal')
    plotGridScalarNodes(triang, z, cmap, vmin=vmin, vmax=vmax, label=label)
    plt.gca().set_xlim(xmin, xmax)
    plt.gca().set_ylim(ymin, ymax)
    noTicks(plt.gca())
    plt.title(title)
    f.tight_layout()
    mapscale(mapx, mapy, maplen, units)
    # plt.show()
    plt.savefig(pngfile, bbox_inches='tight')


def drawGrid(
        triang,
        title,
        pngfile,
        w,
        h,
        xmin,
        xmax,
        ymin,
        ymax,
        c,
        mapx,
        mapy,
        maplen,
        units):
    """ Draw using lines the grid elements"""
    f = plt.figure(
        num=None,
        figsize=(
            w,
            h),
        dpi=100,
        facecolor='w',
        edgecolor='w')
    plt.gca().set_aspect('equal')
    plt.triplot(triang, c, linewidth=0.3)
    plt.gca().set_xlim(xmin, xmax)
    plt.gca().set_ylim(ymin, ymax)
    noTicks(plt.gca())
    plt.title(title)
    f.tight_layout()
    mapscale(mapx, mapy, maplen, units)
    # plt.show()
    plt.savefig(pngfile, bbox_inches='tight')


def generateGridPlots(dataDir, runTag, imgDir):

    hgrid = os.path.join(dataDir, 'hgrid.gr3')
    rough = os.path.join(dataDir, 'rough.gr3')
    drag = os.path.join(dataDir, 'drag.gr3')
    tvd = os.path.join(dataDir, 'tvd.gr3')
    estuary = os.path.join(dataDir, 'estuary.gr3')
    t_nudge = os.path.join(dataDir, 't_nudge.gr3')
    s_nudge = os.path.join(dataDir, 't_nudge.gr3')

    grid = gr.readHGrid(hgrid, True)
    z = grid.nodes[:, 3]

    triang = gridToTriangulation(grid)

# compute areas for ER plots
    a = computeAreas(triang)
    er = np.log10(2.0 * np.sqrt(a) / 3.1415926535897)

    drawScalar(
        triang,
        z,
        'Grid Bathymetry',
        imgDir +
        '/bath-all.png',
        6,
        10.5,
        None,
        None,
        None,
        None,
        0,
        4500,
        plt.get_cmap('Spectral_r'),
        'Bathymetry (m)',
        0.1,
        0.8,
        100000.0,
        'km')
    drawScalar(triang, z, 'Grid Bathymetry - Estuary', imgDir +
               '/bath-estuary.png', 10, 6, 329962, 381353, 271524, 306081, -4,
               40, plt.get_cmap('Spectral_r'),
               'Bathymetry (m)', 0.3, 0.8, 5000.0, 'km')
    drawScalar(
        triang, z, 'Grid Bathymetry - Tidal Freshwater', imgDir +
        '/bath-tidalfresh.png', 8, 8, 377880, 514248, 195546, 299563, -6, 20,
        plt.get_cmap('Spectral_r'),
        'Bathymetry (m)', 0.4, 0.8, 20000.0, 'km')

    drawElements(
        triang,
        er,
        'Equivalent radius',
        imgDir +
        '/er-all.png',
        6,
        10.5,
        None,
        None,
        None,
        None,
        1,
        4,
        plt.cm.jet_r,
        'log10(m^2)',
        0.1,
        0.8,
        100000.0,
        'km')
    drawElements(
        triang,
        er,
        'Equivalent radius - Estuary',
        imgDir +
        '/er-estuary.png',
        10,
        6,
        329962,
        381353,
        271524,
        306081,
        1,
        3,
        plt.cm.jet_r,
        'log10(m^2)',
        0.3,
        0.8,
        5000.0,
        'km')
    drawElements(
        triang,
        er,
        'Equivalent radius - Tidal Freshwater',
        imgDir +
        '/er-tidalfresh.png',
        8,
        8,
        377880,
        514248,
        195546,
        299563,
        1,
        2,
        plt.cm.jet_r,
        'log10(m^2)',
        0.4,
        0.8,
        20000.0,
        'km')

    drawGrid(triang, 'Grid', imgDir + '/grid-all.png', 6, 10.5,
             None, None, None, None, 'k-',
             0.1, 0.8, 100000.0, 'km')
    drawGrid(triang, 'Grid - Estuary', imgDir + '/grid-estuary.png', 10, 6,
             329962, 381353, 271524, 306081, 'k-',
             0.3, 0.8, 5000.0, 'km')
    drawGrid(
        triang,
        'Grid - Tidal Freshwater',
        imgDir +
        '/grid-tidalfresh.png',
        8,
        8,
        377880,
        514248,
        195546,
        299563,
        'k-',
        0.4,
        0.8,
        20000.0,
        'km')

    if os.path.exists(drag):
        grid = gr.readHGrid(drag, True)
        z = grid.nodes[:, 3]
        zmin = min(z)
        zmax = max(z)
        if (zmax - zmin) < 1e-10:
            zmin = zmin - 0.001 * zmin
            zmax = zmax + 0.001 * zmax
            zlab = '''drag.gr3: min = max = %.6f''' % (zmin,)
        else:
            zlab = '''drag.gr3: min = %.6f, max = %.6f''' % (zmin, zmax,)
            zmin = None
            zmax = None
        drawScalar(
            triang,
            z,
            zlab,
            imgDir +
            '/drag.png',
            10,
            6,
            329962,
            381353,
            271524,
            306081,
            zmin,
            zmax,
            plt.cm.jet_r,
            'Drag',
            0.3,
            0.8,
            5000.0,
            'km')

    elif os.path.exists(rough):
        grid = gr.readHGrid(rough, True)
        z = grid.nodes[:, 3]
        zmin = min(z)
        zmax = max(z)
        if (zmax - zmin) < 1e-10:
            zmin = zmin - 0.001 * zmin
            zmax = zmax + 0.001 * zmax
            zlab = '''rough.gr3: min = max = %.6f''' % (zmin,)
        else:
            zlab = '''rough.gr3: min = %.6f, max = %.6f''' % (zmin, zmax,)
            zmin = None
            zmax = None
        drawScalar(
            triang,
            z,
            zlab,
            imgDir +
            '/rough-all.png',
            6,
            10.5,
            None,
            None,
            None,
            None,
            zmin,
            zmax,
            plt.cm.jet_r,
            'Roughness',
            0.1,
            0.8,
            100000.0,
            'km')
        drawScalar(
            triang,
            z,
            zlab,
            imgDir +
            '/rough-estuary.png',
            10,
            6,
            329962,
            381353,
            271524,
            306081,
            zmin,
            zmax,
            plt.cm.jet_r,
            'Roughness',
            0.3,
            0.8,
            5000.0,
            'km')
        drawScalar(
            triang,
            z,
            zlab,
            imgDir +
            '/rough-tidalfresh.png',
            8,
            8,
            377880,
            514248,
            195546,
            299563,
            zmin,
            zmax,
            plt.cm.jet_r,
            'Roughness',
            0.4,
            0.8,
            20000.0,
            'km')

    if os.path.exists(tvd):
        grid = gr.readHGrid(tvd, True)
        z = grid.nodes[:, 3]
        drawScalar(
            triang,
            z,
            'tvd.gr3',
            imgDir +
            '/tvd.png',
            10,
            6,
            329962,
            381353,
            271524,
            306081,
            None,
            None,
            plt.cm.jet_r,
            'TVD',
            0.3,
            0.8,
            5000.0,
            'km')

    if os.path.exists(estuary):
        grid = gr.readHGrid(estuary, True)
        z = grid.nodes[:, 3]
        drawScalar(
            triang,
            z,
            'estuary.gr3',
            imgDir +
            '/estuary.png',
            10,
            6,
            329962,
            381353,
            271524,
            306081,
            None,
            None,
            plt.cm.jet_r,
            'Estuary file for hotstart.in',
            0.3,
            0.8,
            5000.0,
            'km')

    if os.path.exists(s_nudge):
        grid = gr.readHGrid(s_nudge, True)
        z = grid.nodes[:, 3]
        drawScalar(
            triang,
            z,
            's_nudge.gr3',
            imgDir +
            '/s_nudge.png',
            6,
            10.5,
            None,
            None,
            None,
            None,
            None,
            None,
            plt.cm.jet_r,
            'Nudging file for salinity',
            0.1,
            0.8,
            100000.0,
            'km')

    if os.path.exists(t_nudge):
        grid = gr.readHGrid(t_nudge, True)
        z = grid.nodes[:, 3]
        drawScalar(
            triang,
            z,
            't_nudge.gr3',
            imgDir +
            '/t_nudge.png',
            6,
            10.5,
            None,
            None,
            None,
            None,
            None,
            None,
            plt.cm.jet_r,
            'Nudging file temperature',
            0.1,
            0.8,
            100000.0,
            'km')

      # generateCdPlot(dataDir,runTag,imgDir)

#$mlab -r "plotgr3('$1/rough.gr3',[320206 402942], [ 264676 311378], [], 1280, 600, '$1/rough.gr3.png');"
#$mlab -r "plotgr3('$1/estuary.gr3',[320206 402942], [ 264676 311378], [], 1280, 600, '$1/estuary.gr3.png');"
#$mlab -r "plotgr3('$1/t_nudge.gr3',[-53143 557833], [-508041 711159], [], 600, 1024, '$1/t_nudge.gr3.png');"
#$mlab -r "plotgr3('$1/s_nudge.gr3',[-53143 557833], [-508041 711159], [], 600, 1024, '$1/s_nudge.gr3.png');"

    return 0


def parseCommandLine():
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option(
        '-r',
        '--runID',
        action='store',
        type='string',
        dest='runTag',
        help='Run ID, used as a label in post-proc.')
    parser.add_option(
        '-d',
        '--dataDirectory',
        action='store',
        type='string',
        dest='dataDir',
        help='directory where model outputs are stored')
    parser.add_option(
        '-o',
        '--imageDirectory',
        action='store',
        type='string',
        dest='imgDir',
        help='directory where generated images are stored')

    (options, args) = parser.parse_args()
    runTag = options.runTag
    dataDir = options.dataDir
    imgDir = options.imgDir

    if runTag is None:
        parser.print_help()
        parser.error('runTag undefined')
    if not dataDir:
        parser.print_help()
        parser.error('dataDir  undefined')
    if imgDir is None:
        parser.print_help()
        parser.error('imgDir undefined')

    print 'Parsed options:'
    print ' - run ID', runTag
    print ' - dataDir', dataDir
    print ' - output dir', imgDir

    generateGridPlots(dataDir, runTag, imgDir)

if __name__ == "__main__":
    parseCommandLine()
