#!/usr/bin/env python3
#
# (C) Copyright 2015-2015 ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation nor
# does it submit to any jurisdiction.
#

from __future__ import print_function

import os
import argparse
from Magics import macro
from sys import platform

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--format', help='Output file format (pdf, png, gif, svg, ...)', default='pdf')
parser.add_argument('-w', '--wind', help='Plot wind fields', action='store_true')
parser.add_argument('-d', '--default', help='Use MAGICS defaults settings', action='store_true')
parser.add_argument('-r', '--raw', help='Turn off MAGICS automatic scaling', action='store_true')
parser.add_argument('-g', '--grid_shading', help='Use grid_shading', action='store_true')
parser.add_argument('-c', '--cell_shading', help='Use cell_shading', action='store_true')
parser.add_argument('-m', '--marker', help='Use marker', action='store_true')
parser.add_argument('-o', '--output', help='Path to output file', metavar='OUTPUT')
parser.add_argument('-D', '--diff', help='Path to other GRIB', metavar='GRIB')
parser.add_argument('-R', '--reference', help='Path to reference GRIB', metavar='GRIB')

parser.add_argument('-E', '--error', help='Path of other GRIB', metavar='GRIB')
parser.add_argument('-l', '--legend', help='Turn on legend', action='store_true')
parser.add_argument('-n', '--polar_north', help='Polar stereographic north', action='store_true')
parser.add_argument('-s', '--polar_south', help='Polar stereographic south', action='store_true')
parser.add_argument('-e', '--europe', help='Europe', action='store_true')

parser.add_argument('grib', type=str, metavar="GRIB", nargs=1, help="GRIB file to plot")
args = parser.parse_args()

print(args)

tmp = None

EPSILON = 1e-7


def error(x, y):

    if x < EPSILON and y < EPSILON:
        return abs(x - y)

    return abs(x - y) / max(abs(x), abs(y))


def diff(x, y):
    return x - y


mproc = macro.mcont
if args.wind:
    mproc = macro.mwind


scaling = 'on'
contour = mproc(contour_automatic_setting='ecchart',)
grib_extra = {}
contour_extra = {}

if args.legend:
    contour_extra['legend'] = 'on'

if args.diff or args.error:
    scaling = 'off'
    import eccodes

    if args.diff:
        path = args.diff
        proc = diff

    if args.error:
        path = args.error
        proc = error

    with open(args.grib[0], 'rb') as f:
        h1 = eccodes.codes_grib_new_from_file(f)

    with open(path, 'rb') as f:
        h2 = eccodes.codes_grib_new_from_file(f)

    v1 = eccodes.codes_get_values(h1)
    v2 = eccodes.codes_get_values(h2)

    e1 = eccodes.codes_get(h1, "packingError")
    e2 = eccodes.codes_get(h2, "packingError")
    EPSILON = min(e1, e2)

    # Version without numpy

    for i in range(len(v1)):
        v1[i] = proc(v1[i], v2[i])

    eccodes.codes_set_values(h1, v1)

    print('min:', min(v1), 'max:', max(v1))

    # TODO: Use tmpfile
    tmp = "/tmp/plot-diff$$.grib"
    with open(tmp, 'wb') as f:
        eccodes.codes_write(h1, f)

    eccodes.codes_release(h1)
    eccodes.codes_release(h2)

    args.grib = [tmp]

    if args.diff:
        grib_extra = {'grib_interpolation_method': 'nearest'}
        contour = [mproc(contour="off",
                         contour_method="linear",
                         contour_shade="on",
                         contour_shade_technique="grid_shading",
                         contour_shade_min_level=EPSILON,
                         contour_label="off",
                         contour_shade_max_level_colour='rgb(0,0,1)',
                         contour_shade_min_level_colour='rgb(0,0,1)', **contour_extra),
                   mproc(contour="off",
                         contour_shade="on",
                         contour_method="linear",
                         contour_shade_technique="grid_shading",
                         contour_shade_max_level=-EPSILON,
                         contour_label="off",
                         contour_shade_max_level_colour='rgb(1,0,0)',
                         contour_shade_min_level_colour='rgb(1,0,0)', **contour_extra)]
    if args.error:
        grib_extra = {'grib_interpolation_method': 'nearest'}
        contour = mproc(contour="off",
                        contour_shade="on",
                        contour_method="linear",
                        contour_shade_technique="grid_shading",
                        contour_shade_min_level=EPSILON - EPSILON / 2.0,
                        contour_label="off", **contour_extra)


if not args.output:
    view = True
    bn = ''.join(x for x in os.path.basename(args.grib[0]) if x.isalnum() or x in "_-.")  # sanitize input grib basename
    args.output = "/tmp/plot.%s.%d.%s" % (bn, os.getpid(), args.format)
else:
    view = False

base, args.format = os.path.splitext(args.output)

output = macro.output(output_formats=[args.format[1:]],
                      output_name_first_page_number='off',
                      #output_width=8000,
                      output_name=base)

# Setting the coordinates of the geographical area

projection = macro.mmap(subpage_upper_right_longitude=180.00,
                        subpage_upper_right_latitude=90.00,
                        subpage_lower_left_latitude=-90.00,
                        subpage_lower_left_longitude=-180.0,
                        subpage_map_projection='cylindrical')


if args.polar_north:
    projection = macro.mmap(subpage_map_projection="polar_stereographic")

if args.polar_south:
    projection = macro.mmap(subpage_map_projection="polar_stereographic",
                            subpage_map_hemisphere="south")

if args.europe:
    projection = macro.mmap(subpage_upper_right_longitude=65.,
                            subpage_map_projection="polar_stereographic",
                            subpage_map_vertical_longitude=0.,
                            subpage_lower_left_longitude=-37.27,
                            subpage_lower_left_latitude=18.51,
                            subpage_upper_right_latitude=51.28)

foreground = macro.mcoast(map_grid='on')
background = macro.mcoast(map_grid='on',
                          map_grid_colour='tan',
                          map_coastline_land_shade='on',
                          map_coastline_land_shade_colour='cream',
                          map_coastline_colour='tan')


# Define a contour

if args.default:
    contour = mproc(**contour_extra)

if args.raw:
    scaling = 'off'

if args.grid_shading:
    contour = mproc(contour="off", contour_shade="on", contour_shade_technique="grid_shading", contour_method="linear", **contour_extra)
    grib_extra = {'grib_interpolation_method': 'nearest'}

if args.cell_shading:
    contour = mproc(contour="off", contour_shade="on", contour_shade_technique="cell_shading", **contour_extra)

if args.marker:
    contour = mproc(contour="off", contour_shade="on", contour_shade_technique="marker", **contour_extra)

data = []
for i in range(1, 2):
    grib = macro.mgrib(grib_input_file_name=args.grib[0], grib_field_position=i, grib_automatic_scaling=scaling, **grib_extra)
    print(grib)

    data.append(background)
    data.append(grib)
    if isinstance(contour, list):
        for c in contour:
            data.append(c)
    else:
        data.append(contour)
    data.append(foreground)
    data.append(macro.mtext())
    data.append(macro.page())


print('output:', output, 'projection:', projection, 'data:', data)

macro.plot(output, projection, data)

if tmp:
    os.unlink(tmp)

if view:
    op = "open" if platform == "darwin" else "xdg-open"
    os.system("%s %s" % (op, args.output, ))
