#!/usr/bin/env python3
#
# (C) Copyright 1996- 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.

import argparse
import os
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


def miss(x, y):
    return m1 if (x == m1) == (y == m2) else [1, 2][y == m2]


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 = miss
    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.0,
        subpage_map_projection="polar_stereographic",
        subpage_map_vertical_longitude=0.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,))