#!/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_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/" 
    ai_dir = "res_aes_ai/"
    # 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 ]
    ai_jsons =      [ list(map(lambda s: read_json(ai_dir +"{}.res".format(
        s.replace("rijndael_4_4_","128-")
        .replace("rijndael_4_8_","256-")
        .replace("aes_sk_", "0-")
        .replace("_5_1_0","")
    )), l)) for l in inst_ids ]
    # print([ list(map(lambda s:ai_dir +"{}.res".format(
    #     s.replace("rijndael_4_4_","128-")
    #     .replace("rijndael_4_8_","256-")
    #     .replace("_5_1_0","")
    # ), 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 ]
    # 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 ai_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__":
    # ### AES
    xor = 5
    diff = 1
    tr = 0
    inst_names = []
    inst_ids = []
    for version, round_range in [("sk", range(3,4+1)), ("4_4", range(3,5+1)), ("4_6", range(3,10+1)), ("4_8", range(3,14+1))]:
        if version == "4_4":
            vname = "128"
        elif version == "4_6":
            vname = "192"
        elif version == "4_8":
            vname = "256"
        else:
            vname = version
        inst_names += [
            # ["AES\\_{}\\_{}".format(vname, r) for r in round_range]
            # ["{}\\_{}".format(vname, r) for r in round_range]
            ["{}".format(r) for r in round_range]
        ]
        if version == "sk":
            inst_ids += [
                ["aes_{}_{}_{}_{}_{}".format(version, r, xor, diff, tr) for r in round_range]
            ]
        else:
            inst_ids += [
                ["rijndael_{}_{}_{}_{}_{}".format(version, r, xor, diff, tr) for r in round_range]
            ]
    # for t in ["topt", "tenum"]:
    #     insts_curve(inst_names, inst_ids, "aes_{}_{}_{}_{}.pdf".format(t, xor, diff, tr), t, aspect_ratio=2 if t=="tenum" else 1)
    insts_curve_combined(inst_names, inst_ids, "c_aes_{}_{}_{}.pdf".format(xor, diff, tr), ["SK", "RK AES 128", "RK AES 192", "RK AES 256"])