#!/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
import tempfile
from sys import platform

from Magics import macro

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('-o', '--output', help='Path to output file', metavar='OUTPUT')
parser.add_argument('-e', '--europe', help='Europe', action='store_true')

parser.add_argument('--land', help='Map coastline land shade colour', default='cream')
parser.add_argument('--sea', help='Map coastline sea shade colour', default='')
parser.add_argument('--width', help='Plot width', type=int, default=0)

g = parser.add_mutually_exclusive_group()
g.add_argument('-n', '--polar_north', help='Polar stereographic North', action='store_true')
g.add_argument('-s', '--polar_south', help='Polar stereographic South', action='store_true')

g = parser.add_mutually_exclusive_group()
g.add_argument('-g', '--grid_shading', help='Use grid_shading', action='store_true')
g.add_argument('-c', '--cell_shading', help='Use cell_shading', action='store_true')
g.add_argument('-m', '--marker', help='Use marker', action='store_true')

g = parser.add_mutually_exclusive_group()
g.add_argument('-D', '--diff', help='Difference to another GRIB', metavar='GRIB')
g.add_argument('-E', '--error', help='Error to another GRIB', metavar='GRIB')
g.add_argument('-M', '--miss', help='Compare missing values to another GRIB', metavar='GRIB')

def add_feature(name, **kwargs):
    g = parser.add_mutually_exclusive_group()
    g.add_argument('--{}'.format(name), dest=name, action='store_true')
    g.add_argument('--no-{}'.format(name), dest=name, action='store_false') 
    parser.set_defaults(**kwargs)

add_feature('grid', grid=True)
add_feature('legend', legend=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 = {}
output_extra = {}

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

if args.width:
    output_extra['output_width'] = args.width

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

    path = args.diff if args.diff else args.error if args.error else args.miss

    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)

    if args.miss:
        b1 = eccodes.codes_get(h1, "bitmapPresent")
        b2 = eccodes.codes_get(h2, "bitmapPresent")
        assert b1 and b2, "both files must have missing values"

        m1 = eccodes.codes_get(h1, "missingValue")
        m2 = eccodes.codes_get(h2, "missingValue")
        proc = lambda x, y: m1 if (x == m1) == (y == m2) else [1,2][y == m2]
    elif args.diff :
        proc = diff
    else:
        proc = error

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

    tmp = tempfile.NamedTemporaryFile(mode='wb')
    args.grib = [tmp.name]

    eccodes.codes_write(h1, tmp.file)

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

    grib_extra = {'grib_interpolation_method': 'nearest'}
    if args.diff:
        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:
        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 args.miss:
        contour = mproc(contour="off",
                        contour_max_level=2.1,
                        contour_level_count=2,
                        contour_level_tolerance=0,
                        contour_shade="on",
                        contour_shade_technique="marker",
                        contour_shade_colour_method="list",
                        contour_shade_colour_list=["red","blue"],
                        **contour_extra)


if not args.output:
    view = True
    with tempfile.NamedTemporaryFile(mode='wb', suffix='.' + args.format, delete=False) as f:
        args.output = f.name
        print('args.output:', args.output)
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_name=base,
                      **output_extra)

# 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=args.grid)
background = macro.mcoast(map_grid=args.grid,
                          map_grid_colour='tan',
                          map_coastline_land_shade=bool(args.land),
                          map_coastline_land_shade_colour=args.land,
                          map_coastline_sea_shade=bool(args.sea),
                          map_coastline_sea_shade_colour=args.sea,
                          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 view:
    op = "open" if platform == "darwin" else "xdg-open"
    os.system("%s %s" % (op, args.output, ))
