#!/usr/bin/python3
#***************************** LICENSE START ***********************************
#
# Copyright 2018 ECMWF and INPE. This software is distributed under the terms
# of the Apache License version 2.0. In applying this license, ECMWF does not
# waive the privileges and immunities granted to it by virtue of its status as
# an Intergovernmental Organization or submit itself to any jurisdiction.
#
# ***************************** LICENSE END *************************************

import json
import sys
import os
import datetime

script_command_prefix = {"macro": "","python": "mv."}
script_param_assign_char = {"macro": ":","python": "="}
script_translate_class_param = { "macro": "class", "python": "class_"}
script_file_suffix = {"macro": ".mv","python": ".py"}


def load_json_mars(mars_file):
    """ Reads json file and returns dictionary """
    with open(mars_file, 'r') as defs:
        dictionary_of_defs = json.load(defs)
    return dictionary_of_defs


def load_json_layer(layer_file):
    """ Reads json file and returns dictionary """

    with open(layer_file, 'r') as defs:
        dictionary_of_defs = json.load(defs)
    return dictionary_of_defs


def build_title_request(main_title, short_name):
    """ Generate Metview request for title """

    start_text = "<grib_info key=\\'base-date\\' format=\\'%d.%m.%Y. %H\\' where=\\'shortName=" \
                 + short_name + "\\' /> UTC"
    step_text = "<grib_info key=\\'step\\' where=\\'shortName=" + short_name + "\\' />h"
    valid_text = "<grib_info key=\\'valid-date\\' format=\\'%d.%m.%Y. %H\\' where=\\'shortName=" + short_name \
                 + "\\' /> UTC"

    title = "MTEXT,\n" \
        + " TEXT_LINE_1 = '" + main_title + "  START: " + start_text + "   STEP: " + step_text \
        + "  VALID FOR: " + valid_text + "',\n" \
        + " TEXT_COLOUR = CHARCOAL"

    return title


def build_title_script(main_title, short_name, mode):
    """ Generate script code for title """

    start_text = "<grib_info key='base-date' format='%d.%m.%Y. %H' where='shortName=" \
                 + short_name + "' /> UTC"
    step_text = "<grib_info key='step' where='shortName=" + short_name + "' />h"
    valid_text = "<grib_info key='valid-date' format='%d.%m.%Y. %H' where='shortName=" + short_name \
                 + "' /> UTC"

    sep = script_param_assign_char[mode]

    title = script_command_prefix[mode] + "mtext(\n" \
        + "    text_line_1 " + sep + " \"" + main_title + "  START: " + start_text + "  STEP:  " + step_text \
        + "  VALID FOR =  " + valid_text + "\",\n" \
        + "    text_colour " + sep + " \"charcoal\")"

    return title


def build_visdef_request(style_name):
    """ Generate Metview request for contouring """

    visdef = "MCONT,\n" \
        + " CONTOUR_AUTOMATIC_SETTING = STYLE_NAME, \n" \
        + " CONTOUR_STYLE_NAME = " + style_name + ",\n" \
        + " LEGEND = ON"

    return visdef


def build_visdef_script(style_name, mode):
    """ Generate script code for contouring """

    sep = script_param_assign_char[mode]

    visdef = script_command_prefix[mode] + "mcont(\n" \
        + "    contour_automatic_setting " + sep + " \"style_name\",\n" \
        + "    contour_style_name " + sep + "        \"" + style_name + "\",\n" \
        + "    legend " +  sep + "                   \"on\")"

    return visdef


def build_legend_request():
    """ Generate Metview request for legend """

    legend_req = "MLEGEND,\n" \
        + " LEGEND_TEXT_COLOUR = CHARCOAL"

    return legend_req


def build_legend_script(mode):
    """ Generate script code for legend """

    sep = script_param_assign_char[mode]

    legend = script_command_prefix[mode] + "mlegend(\n" \
        + "    legend_text_colour " + sep +" \"charcoal\")"

    return legend


def build_mars_script(mars, mode):
    """ Generate code for MARS request """

    sep = script_param_assign_char[mode]

    text_for_mars = script_command_prefix[mode] + 'retrieve(\n'
    for key in mars:
        v = mars[key]

        if type(v) == str:

            if not v.startswith('[') or not v.endswith(']'):
                v = "'" + v + "'"

        if key == "class" :
            key = script_translate_class_param[mode]

        text_for_mars += '    ' + str(key) + ' ' + sep + ' ' + str(v) + ',\n'

    text_for_mars = text_for_mars[:-2]
    text_for_mars += ')'
    return text_for_mars


def build_pproc_script_call(mars, pproc_func, short_name, pproc_func_extra_args, mode):
    """ Generate code for calling post-processing of retrieved fields """

    params = mars["param"]
    if type(params) == list:
        params_str = ", ".join(['\"' + s + '\"' for s in params])
        # params_str = ", ".join(params)
        params_str = '[' + params_str + ']'
    else:
        params_str = '[' + params + ']'

    text_for_mars = pproc_func + '(' + params_str + ', data , \"' + short_name + "\""
    for a in pproc_func_extra_args:
        text_for_mars += ', ' + a

    text_for_mars += ')'

    return text_for_mars


def build_pproc_script_code(pproc_func, mode):
    """ Generate code for post-processing of retrieved fields """

    suffix = script_file_suffix[mode]
    filename = os.environ["METVIEW_DIR_SHARE"] + "/eccharts/scripts/" + pproc_func + suffix
    with open(filename, 'r') as fp:
        return fp.read()


def steps_for_interval(step_str, interval, layer_name):

    # define the available steps in MARS
    if layer_name in {'mn2t', 'mx2t', '10m_fg_interval'}:
        all_steps = list(range(0, 91))
        all_steps.extend(list(range(93, 147, 3)))
        all_steps.extend(list(range(150, 246, 6)))
    else:
        return ""

    # convert the steps into an integer list
    if not step_str.startswith('[') or not step_str.endswith(']'):
        print("\'step\' is not formatted as a list:", step_str)
        sys.exit(1)

    steps = step_str[1:-1].split(",")
    print("steps=", steps)
    steps = [int(i) for i in steps]
    print("steps=", steps)

    # define the steps for the given interval
    r_steps = []
    for ts in steps:

        step_start = ts - interval
        step_end = ts

        try:
            idx_start = all_steps.index(step_start)
        except ValueError:
            print("Cannot retrieve data for interval ending at step:", ts, "h")
            raise

        try:
            idx_end = all_steps.index(step_end)
        except ValueError:
            print("Cannot retrieve step:", step_end, "h for interval ending at step:", ts, "h")
            raise

        r_steps.extend(all_steps[idx_start+1:idx_end+1])

    # print("r_steps=", r_steps)

    # remove duplicates and sort
    r_steps = sorted(list(set(r_steps)), key=int)

    # return as a string
    return "[" + ",".join([str(i) for i in r_steps]) + "]"


def main():
    
    # get arguments
    if len(sys.argv) < 8:
        print("Too few arguments = " + str(len(sys.argv)) + ". Minimum number of arguments is 8.")
        sys.exit(1)

    mode = sys.argv[1]
    if mode not in ("data", "macro", "python"):
        print("Incorrect first argument: " + mode + " Allowed values: \"data\",\"macro\" or \"python\"")
        sys.exit(1)

    layer_def_file = sys.argv[2]
    mars_def_file = sys.argv[3]
    script_file = sys.argv[4]
    request_file = sys.argv[5]
    layer_name = sys.argv[6]
    style_name = sys.argv[7]
    need_title = sys.argv[8]
    if need_title == "1":
        need_title = True
    else:
        need_title = False

    # Read json defs
    layer_def = load_json_layer(layer_def_file)
    mars_def = load_json_mars(mars_def_file)

    # -----------------------------------
    # Build macro code for MARS retrieval
    # -----------------------------------

    extra_mars_options = {}
    interval = ""

    # Read arguments for the mars request
    arguments = sys.argv
    for v in arguments[9:]:
        try:
            (key, val) = v.split('=')
        except ValueError:
            print("No key is defined in argument: ", v)
            raise

        # 'interval' is not a mars key. If it is defined we need to modify
        #  the steps so that a proper interval could be retrieved
        if key != 'interval':
            extra_mars_options[key] = val
        else:
            interval = val

    # Get MARS request from layer definition
    mars = mars_def[layer_name]['mars']

    # Update it with the extra options
    mars.update(extra_mars_options)

    if interval:
        steps = mars['step']
        modified_steps = steps_for_interval(mars['step'], int(interval), layer_name)
        if modified_steps:
            mars['step'] = modified_steps
        else:
            print("Could not define steps for the specified interval=", interval)
            sys.exit(1)

    # print("name=", layer_name)
    # print("title=", type(layer_def))
    # print("mars=", mars)

    # -----------------------------------
    # Find out title and style
    # -----------------------------------

    # Get layer title
    layer_title = layer_def[layer_name]['title']

    # Get parameter shortName
    short_name = mars_def[layer_name]['shortName']

    # Get parameter post-processing option
    pproc_func = mars_def[layer_name].get('proc', "")

    # for certain layers we need to add extra parameters to the post processing
    # script
    pproc_func_extra_args = []

    if interval:
        pproc_func_extra_args.append(steps)
        pproc_func_extra_args.append(interval)

    # Get default style name if style is not defined
    if style_name == "":
        style_name = layer_def[layer_name]['style']

    # -----------------------------------
    # In data mode:
    # -generate macro to retrieve data
    # -define the title, contouring and legend as Metview requests
    # -----------------------------------

    if mode == "data":

        # -----------------------------------
        # Macro code to retrieve the data
        # -----------------------------------

        macro_text = "#Metview macro\n\n" \
                + "args=arguments()\n" \
                + "print(args)\n" \
                + "outputFile=args[1]\n" \
                + "data = " + build_mars_script(mars, "macro") + "\n\n"

        if pproc_func:
            macro_text += "data = " + build_pproc_script_call(mars, pproc_func, short_name, pproc_func_extra_args, "macro") + "\n\n"

        macro_text += "write(outputFile,data)\n\n" \
                + "return(0)"

        if pproc_func:
            macro_text += "\n\n" + build_pproc_script_code(pproc_func, "macro")

        req_text = ''

        # --------------------------------
        # Build title request
        # --------------------------------
        if need_title:
            req_text += build_title_request(layer_title, short_name)

        # --------------------------------
        # Build visdef request
        # --------------------------------

        req_text += "\n\n" + build_visdef_request(style_name)

        # --------------------------------
        # Build legend request
        # --------------------------------

        req_text += "\n\n" + build_legend_request() + "\n"

        # --------------------------------
        # Save request into file
        # --------------------------------
        # print(req_text)

        with open(request_file, 'w') as fp:
            fp.write(req_text)

        # -----------------------
        #  Save macro into file
        # -----------------------
        # print(macro_text)

        with open(script_file, 'w') as fp:
            fp.write(macro_text)

    # -----------------------------------------------------------------------------
    # In macro or python mode mode generate script to retrieve and plot the data
    # -----------------------------------------------------------------------------

    else:

        if mode == "macro" :
            script_text = "# Metview Macro\n\n"
        else:
            script_text = "import metview as mv\n\n"

        # --------------------------------
        #  Add licence text
        # --------------------------------

        with open(os.environ["METVIEW_DIR_SHARE"] + "/etc/licence_for_macros.txt", 'r') as fp:
            licence_text = fp.read().replace("[YEAR]", str(datetime.datetime.now().year))
            for s in licence_text.split("\n"):
                script_text += "# " + s + "\n"

        # -------------------------------------
        # Build post-processing code - Python
        # -------------------------------------

        if mode == "python" and pproc_func:
            script_text += "\n# The post-processing function\n" + build_pproc_script_code(pproc_func, mode) + "\n\n"

        # --------------------------------
        # Build mars retrieval
        # --------------------------------

        script_text += "\n# Retrieve data from MARS\ndata = " + build_mars_script(mars, mode) + "\n\n"

        # --------------------------------
        # Build post-processing call
        # --------------------------------

        if pproc_func:
            script_text += "# Perform post-processing on data\ndata = " \
                    + build_pproc_script_call(mars, pproc_func, short_name, pproc_func_extra_args, mode) + "\n\n"

        # --------------------------------
        # Build title
        # --------------------------------
        if need_title:
            script_text += "# Define title\ntitle = " + build_title_script(layer_title, short_name, mode) + "\n\n"

        # --------------------------------
        # Build visdef
        # --------------------------------

        script_text += "# Define contouring\ncont = " + build_visdef_script(style_name, mode) + "\n\n"

        # --------------------------------
        # Add legend
        # --------------------------------

        script_text += "# Define legend\nlegend = " + build_legend_script(mode) + "\n\n"

        # ---------------------------------
        # Plot command
        # ---------------------------------

        script_text += "# Generate plot\n" + script_command_prefix[mode]
        if need_title:
            script_text += "plot(data, title, cont, legend)\n\n"
        else:
            script_text += "plot(data, cont, legend)\n\n"

        # -----------------------------------
        # Build post-processing code - Macro
        # -----------------------------------

        if mode == "macro" and pproc_func:
            script_text += "\n# The post-processing function\n" + build_pproc_script_code(pproc_func, mode) + "\n\n"

        # -----------------------
        #  Save macro into file
        # -----------------------

        with open(script_file, 'w') as fp:
            fp.write(script_text)


if __name__ == "__main__":    
    main()