#!/usr/bin/python3
from helper import read_json

import matplotlib
import matplotlib.pyplot as plt

def flatten(l):
    res = []
    for e in l:
        res += e
    return res


def gen_plot(instances, algos, ax):
    """
    instances: 2d array with instance names
    algos: [{name:String, style:String, values:2d array for each instance class, for each instance, value or None}, ...]
    """
    matplotlib.rcParams['text.usetex'] = True
    matplotlib.rcParams.update({
        "pgf.texsystem": "pdflatex",
        'font.family': 'serif',
        'text.usetex': True,
        'pgf.rcfonts': False,
    })
    SMALL_SIZE = 8
    matplotlib.rc('font', size=SMALL_SIZE)
    matplotlib.rc('axes', titlesize=SMALL_SIZE) 
    matplotlib.rcParams['lines.markersize'] = 3.
    inst_list = flatten(instances)
    inst_range = list(range(0,len(inst_list)))
    # draw instance class separators
    xi = 0
    for e in instances:
        ax.axvline(x=xi, color="lightgray", linewidth=2, linestyle="-", label=None)
        xi += len(e)
    # draw algo values
    for algo in algos:
        name, style, values = algo["name"], algo["style"], algo["values"]
        xi = 0
        for (subsequence_pos,value_class) in enumerate(values):
            x = []
            y = []
            for v in value_class:
                if v != None:
                    x.append(xi)
                    y.append(v)
                xi += 1
            ax.plot(x, y, style, linewidth=0.5, label=(name if subsequence_pos == 0 else None))
    # other draw instructions
    plt.xticks(inst_range, inst_list, rotation=45, ha="right", fontsize=7, rotation_mode="anchor")
    ax.grid(True, color="lightgray", linestyle="--")
    # plt.xlim([-0.2, len(inst_range)])



# def insts_curve(inst_names, inst_ids, output_filename, value, aspect_ratio=2):
#     """
#     inst_names: vector of instance names (to be displayed)
#     inst_ids: identifier for a given instance
#     output_filename: pdf file to produce
#     value: type of value [topt, tenum]
#     """
#     picat_dir =   "res_picat/"  
#     chuffed_dir = "res_chuffed/"
#     gurobi_dir =  "res_gurobi/" 
#     # gecode_dir = "res_gecode/"
#     picat_jsons =   [ list(map(lambda s: read_json(picat_dir  +"{}.json.mzn".format(s)), l)) for l in inst_ids ]
#     chuffed_jsons = [ list(map(lambda s: read_json(chuffed_dir+"{}.json.mzn".format(s)), l)) for l in inst_ids ]
#     gurobi_jsons =  [ list(map(lambda s: read_json(gurobi_dir +"{}.json.mzn".format(s)), l)) for l in inst_ids ]
#     # gecode_jsons =  [ list(map(lambda s: read_json(gecode_dir +"{}.json.mzn".format(s)), l)) for l in inst_ids ]
#     plt = plt.figure()
#     gen_plot(inst_names, [
#         {"name": "Picat SAT", "style": "bo-", "values": [
#             list(map(lambda e: e[value] if value in e else None, l)) for l in picat_jsons
#         ]},
#         {"name": "Chuffed", "style": "g+-", "values": [
#             list(map(lambda e: e[value] if value in e else None, l)) for l in chuffed_jsons
#         ]},
#         {"name": "Gurobi", "style": "rx-", "values": [
#             list(map(lambda e: e[value] if value in e else None, l)) for l in gurobi_jsons
#         ]}#,
#         # {"name": "Gecode", "style": "mv-", "values": [
#         #     list(map(lambda e: e[value] if value in e else None, l)) for l in gecode_jsons
#         # ]}
#     ], plt)
#     plt.ylabel("Time in seconds (lower is better)")
#     plt.yscale("log")
#     # plt.xlabel(r"Instances")
#     plt.legend(loc="lower right")
#     plt.title(r"Solvers performance ({}). Time limit: 1 hour".format(
#         "$t_{opt}$" if value == "topt" else "$t_{enum}$"
#     ))
#     plt.gca().set_aspect(aspect_ratio)
#     plt.savefig(output_filename, bbox_inches='tight')
#     plt.close()



def insts_curve_combined(inst_names, inst_ids, output_filename, group_labels=[]):
    """
    inst_names: vector of instance names (to be displayed)
    inst_ids: identifier for a given instance
    output_filename: pdf file to produce
    ratio_opt: aspect ratio for opt
    ratio_enum: aspect ratio for enum
    group_labels: by default, empty, otherwise, indicates to the top of the plot the name of the category
    """
    # directories where to find results
    picat_dir =   "res_picat/"  
    chuffed_dir = "res_chuffed/"
    gurobi_dir =  "res_gurobi/"
    soa_dir = "res_soa_midori/"
    # gecode_dir = "res_gecode/"
    # list of json files using the instance list given in parameters
    picat_jsons =   [ list(map(lambda s: read_json(picat_dir  +"{}.json.mzn".format(s)), l)) for l in inst_ids ]
    chuffed_jsons = [ list(map(lambda s: read_json(chuffed_dir+"{}.json.mzn".format(s)), l)) for l in inst_ids ]
    gurobi_jsons =  [ list(map(lambda s: read_json(gurobi_dir +"{}.json.mzn".format(s)), l)) for l in inst_ids ]
    soa_jsons =      [ list(map(lambda s: read_json(soa_dir +"{}.res".format(
        s.replace("midori_64_","64-")
        .replace("midori_128_","128-")
        .replace("midori_sk_", "0-")
        .replace("_4_1_0","")
    )), l)) for l in inst_ids ]
    # create a figure (with size ratio of 10/3)
    fig = plt.figure(figsize=(10, 3))
    # separate the figure into two parts with no space between them
    gs = fig.add_gridspec(2, hspace=0)
    # use the same axis for x and y
    axs = gs.subplots(sharex=True,sharey=True)
    # for each subfigure (topt, tenum)
    for (i,value) in enumerate(["topt", "tenum"]):
        # call function that extracts the result and plots the curve in axs[i]
        gen_plot(inst_names, [
            {"name": "Picat SAT", "style": "bo-", "values": [
                list(map(lambda e: e[value] if value in e else None, l)) for l in picat_jsons
            ]},
            {"name": "Chuffed", "style": "g+-", "values": [
                list(map(lambda e: e[value] if value in e else None, l)) for l in chuffed_jsons
            ]},
            {"name": "Gurobi", "style": "rx-", "values": [
                list(map(lambda e: e[value] if value in e else None, l)) for l in gurobi_jsons
            ]},
            {"name": "state-of-the-art", "style": "k--", "values": [
                list(map(lambda e: e[value] if value in e else None, l)) for l in soa_jsons
            ]}
        ], axs[i])
    # set labels, scales, and legend
    axs[0].set_ylabel("$t_{opt}$ (s)")
    axs[1].set_ylabel("$t_{enum}$ (s)")
    plt.yscale("log")
    plt.legend(loc="lower right")
    # instance group names
    if group_labels != []:
        group_xmin = 0
        group_xmax = 0
        for i,e in enumerate(inst_names):
            group_xmin = group_xmax
            group_xmax += len(e)
            x_pos = (group_xmin+group_xmax)/2.
            axs[0].text(x_pos,7200,group_labels[i], ha="center")
    # export to pdf
    plt.savefig(output_filename, bbox_inches='tight')
    plt.close()



if __name__ == "__main__":
    ### MIDORI
    xor = 4
    diff = 1
    tr = 0
    inst_names = []
    inst_ids = []
    for version, round_range in [("sk", range(3,5+1)), (64, range(3,16+1)), (128, range(3,20+1))]:
        inst_names += [
            # ["Midori\\_{}\\_{}".format(version, r) for r in round_range]
            ["{}".format(r) for r in round_range]
        ]
        inst_ids += [
            ["midori_{}_{}_{}_{}_{}".format(version, r, xor, diff, tr) for r in round_range]
        ]
    insts_curve_combined(
        inst_names,
        inst_ids,
        "c_midori_{}_{}_{}.pdf".format(xor, diff, tr),
        ["SK", "Midori-64 RK", "Midori-128 RK"]
    )