#!/usr/bin/python3
from helper import read_json, extract_labels, extract_v
from sys import argv

import matplotlib
import matplotlib.pyplot as plt

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


def gen_plot(instances, algos):
    """
    instances: 2d array with instance names
    algos: [{name:String, style:String, values:2d array for each instance class, for each instance, value or None}, ...]
    """
    plt.figure()
    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:
        plt.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
            plt.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")
    plt.grid(True, color="lightgray", linestyle="--")
    plt.xlim([-0.2, len(inst_range)])



def aes_curve(diff, tr, output_filename, title, value):
    """
    diff: use diff or not [0,1]
    tr: use transitivity constraints or not [0,1]
    output_filename: pdf file to produce
    title: figure name
    value: type of value [topt, tenum]
    """
    picat_dir =   "21_03_05_rijndael_picat/"
    chuffed_dir = "21_03_15_rijndael_chuffed/"
    gurobi_dir =  "21_03_15_rijndael_gurobi/"
    ###################### AES CURVES ######################################
    range1 = range(3,5+1)
    range2 = range(3,10+1)
    range3 = range(3,14+1)
    inst_names = [
        ["AES\\_128\\_{}".format(r) for r in range1],
        ["AES\\_192\\_{}".format(r) for r in range2],
        ["AES\\_256\\_{}".format(r) for r in range3]
    ]
    picat = [
        ["rijndael_{}_4_4_{}_{}".format(r, diff, tr) for r in range1],
        ["rijndael_{}_4_6_{}_{}".format(r, diff, tr) for r in range2],
        ["rijndael_{}_4_8_{}_{}".format(r, diff, tr) for r in range3]
    ]
    chuffed = [
        ["rijndael_{}_4_4_{}_{}_0".format(r, diff, tr) for r in range1],
        ["rijndael_{}_4_6_{}_{}_0".format(r, diff, tr) for r in range2],
        ["rijndael_{}_4_8_{}_{}_0".format(r, diff, tr) for r in range3]
    ]
    gurobi = [
        ["rijndael_{}_4_4_{}_{}_1".format(r, diff, tr) for r in range1],
        ["rijndael_{}_4_6_{}_{}_1".format(r, diff, tr) for r in range2],
        ["rijndael_{}_4_8_{}_{}_1".format(r, diff, tr) for r in range3]
    ]
    picat_jsons =   [ list(map(lambda s: read_json(picat_dir  +"{}.mzn".format(s)), l)) for l in picat   ]
    chuffed_jsons = [ list(map(lambda s: read_json(chuffed_dir+"{}.mzn".format(s)), l)) for l in chuffed ]
    gurobi_jsons =  [ list(map(lambda s: read_json(gurobi_dir +"{}.mzn".format(s)), l)) for l in gurobi  ]
    gen_plot(inst_names, [
        {"name": "picat", "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
        ]}
    ])
    plt.ylabel("time in seconds (lower is better)")
    plt.yscale("log")
    plt.xlabel(r"AES instances")
    plt.legend(loc="lower right")
    plt.title(title)
    plt.gca().set_aspect(2)
    plt.savefig(output_filename, bbox_inches='tight')
    plt.close()




def midori_curve(diff, tr, output_filename, title, value, version):
    """
    diff: use diff or not [0,1]
    tr: use transitivity constraints or not [0,1]
    output_filename: pdf file to produce
    title: figure name
    value: type of value [topt, tenum]
    version: midori version [64,128]
    """
    picat_dir =   "21_03_05_midori{}_picat/".format(version)
    chuffed_dir = "21_03_05_midori{}_chuffed/".format(version)
    gurobi_dir =  "21_03_05_midori{}_gurobi/".format(version)
    ###################### AES CURVES ######################################
    rmax = 16
    if version == 128:
        rmax = 20
    range1 = range(3,rmax+1)
    inst_names = [
        ["Midori\\_{}\\_{}".format(version, r) for r in range1]
    ]
    picat = [
        ["midori{}_{}_{}_{}".format(version, r, diff, tr) for r in range1]
    ]
    chuffed = [
        ["midori{}_{}_{}_{}".format(version, r, diff, tr) for r in range1]
    ]
    gurobi = [
        ["midori{}_{}_{}_{}_1".format(version, r, diff, tr) for r in range1]
    ]
    picat_jsons =   [ list(map(lambda s: read_json(picat_dir  +"{}.mzn".format(s)), l)) for l in picat   ]
    chuffed_jsons = [ list(map(lambda s: read_json(chuffed_dir+"{}.mzn".format(s)), l)) for l in chuffed ]
    gurobi_jsons =  [ list(map(lambda s: read_json(gurobi_dir +"{}.mzn".format(s)), l)) for l in gurobi  ]
    gen_plot(inst_names, [
        {"name": "picat", "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
        ]}
    ])
    plt.ylabel("time in seconds (lower is better)")
    plt.yscale("log")
    plt.xlabel(r"Midori instances")
    plt.legend(loc="lower right")
    plt.title(title)
    plt.gca().set_aspect(2)
    plt.savefig(output_filename, bbox_inches='tight')
    plt.close()


rijndael_rounds = {
    (4,4): 6,
    (4,5): 8,
    (4,6): 10,
    (4,7): 12,
    (4,8): 14,
    (5,4): 4,
    (5,5): 6,
    (5,6): 9,
    (5,7): 10,
    (5,8): 12,
    (6,4): 3,
    (6,5): 5,
    (6,6): 7,
    (6,7): 9,
    (6,8): 10,
    (7,4): 3,
    (7,5): 4,
    (7,6): 6,
    (7,7): 7,
    (7,8): 9,
    (8,4): 3,
    (8,5): 4,
    (8,6): 5,
    (8,7): 6,
    (8,8): 8
}

def rijndael_curve(diff, tr, output_filename, title, value, x):
    """
    diff: use diff or not [0,1]
    tr: use transitivity constraints or not [0,1]
    output_filename: pdf file to produce
    title: figure name
    value: type of value [topt, tenum]
    x: {4..8}
    """
    picat_dir =   "21_03_05_rijndael_picat/"
    chuffed_dir = "21_03_15_rijndael_chuffed/"
    gurobi_dir =  "21_03_15_rijndael_gurobi/"
    ###################### RIJNDAEL CURVES ######################################
    ranges = [range(3, rijndael_rounds[(x,k)]+1) for k in range(4,8+1)]
    inst_names = [
        ["Rijndael\\_{},{}\\_{}".format(x,i+4,r) for r in ranges_elt] for (i,ranges_elt) in enumerate(ranges)
    ]
    picat = [
        ["rijndael_{}_{}_{}_{}_{}".format(r,x,i+4,diff,tr) for r in ranges_elt] for (i,ranges_elt) in enumerate(ranges)
    ]
    chuffed = [
        ["rijndael_{}_{}_{}_{}_{}_0".format(r,x,i+4,diff,tr) for r in ranges_elt] for (i,ranges_elt) in enumerate(ranges)
    ]
    gurobi = [
        ["rijndael_{}_{}_{}_{}_{}_1".format(r,x,i+4,diff,tr) for r in ranges_elt] for (i,ranges_elt) in enumerate(ranges)
    ]
    picat_jsons =   [ list(map(lambda s: read_json(picat_dir  +"{}.mzn".format(s)), l)) for l in picat   ]
    chuffed_jsons = [ list(map(lambda s: read_json(chuffed_dir+"{}.mzn".format(s)), l)) for l in chuffed ]
    gurobi_jsons =  [ list(map(lambda s: read_json(gurobi_dir +"{}.mzn".format(s)), l)) for l in gurobi  ]
    gen_plot(inst_names, [
        {"name": "picat", "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
        ]}
    ])
    plt.ylabel("time in seconds (lower is better)")
    plt.yscale("log")
    plt.xlabel(r"Rijndael instances")
    plt.legend(loc="lower right")
    plt.title(title)
    plt.gca().set_aspect(2)
    plt.savefig(output_filename, bbox_inches='tight')
    plt.close()





if __name__ == "__main__":
    # rijndael curves
    for x in range(4,8+1):
        # topt
        rijndael_curve(0, 0, "rijndael_{}_topt_0_0.pdf".format(x),
            r"Solver performance ($t_{opt}$) without diff variables. Time limit: 1 hour",
            "topt", x
        )
        rijndael_curve(1, 0, "rijndael_{}_topt_1_0.pdf".format(x),
            r"Solver performance ($t_{opt}$) with diff variables. Time limit: 1 hour",
            "topt", x
        )
        rijndael_curve(1, 1, "rijndael_{}_topt_1_1.pdf".format(x),
            r"Solver performance ($t_{opt}$) with diff variables and transitivity constraints. Time limit: 1 hour",
            "topt", x
        )
        # tenum
        rijndael_curve(0, 0, "rijndael_{}_tenum_0_0.pdf".format(x),
            r"Solver performance ($t_{enum}$) without diff variables. Time limit: 1 hour",
            "tenum", x
        )
        rijndael_curve(1, 0, "rijndael_{}_tenum_1_0.pdf".format(x),
            r"Solver performance ($t_{enum}$) with diff variables. Time limit: 1 hour",
            "tenum", x
        )
        rijndael_curve(1, 1, "rijndael_{}_tenum_1_1.pdf".format(x),
            r"Solver performance ($t_{enum}$) with diff variables and transitivity constraints. Time limit: 1 hour",
            "tenum", x
        )
    # AES topt
    aes_curve(
        0,0,
        "aes_topt_0_0.pdf",
        r"Solver performance ($t_{opt}$) without diff variables. Time limit: 1 hour",
        "topt"
    )
    aes_curve(
        1,0,
        "aes_topt_1_0.pdf",
        r"Solver performance ($t_{opt}$) with diff variables. Time limit: 1 hour",
        "topt"
    )
    aes_curve(
        1,1,
        "aes_topt_1_1.pdf",
        r"Solver performance ($t_{opt}$) with diff variables and transitivity constraints. Time limit: 1 hour",
        "topt"
    )
    # AES tenum
    aes_curve(
        0,0,
        "aes_tenum_0_0.pdf",
        r"Solver performance ($t_{enum}$) without diff variables. Time limit: 1 hour",
        "tenum"
    )
    aes_curve(
        1,0,
        "aes_tenum_1_0.pdf",
        r"Solver performance ($t_{enum}$) with diff variables. Time limit: 1 hour",
        "tenum"
    )
    aes_curve(
        1,1,
        "aes_tenum_1_1.pdf",
        r"Solver performance ($t_{enum}$) with diff variables and transitivity constraints. Time limit: 1 hour",
        "tenum"
    )
    # midori64 topt
    midori_curve(
        0, 0, "midori_64_topt_0_0.pdf",
        r"Solver performance ($t_{opt}$) without diff variables. Time limit: 1 hour",
        "topt", 64
    )
    midori_curve(
        1, 0, "midori_64_topt_1_0.pdf",
        r"Solver performance ($t_{opt}$) with diff variables. Time limit: 1 hour",
        "topt", 64
    )
    # midori64 tenum
    midori_curve(
        0, 0, "midori_64_tenum_0_0.pdf",
        r"Solver performance ($t_{enum}$) without diff variables. Time limit: 1 hour",
        "tenum", 64
    )
    midori_curve(
        1, 0, "midori_64_tenum_1_0.pdf",
        r"Solver performance ($t_{enum}$) with diff variables. Time limit: 1 hour",
        "tenum", 64
    )
    # midori128 topt
    midori_curve(
        0, 0, "midori_128_topt_0_0.pdf",
        r"Solver performance ($t_{opt}$) without diff variables. Time limit: 1 hour",
        "topt", 128
    )
    midori_curve(
        1, 0, "midori_128_topt_1_0.pdf",
        r"Solver performance ($t_{opt}$) with diff variables. Time limit: 1 hour",
        "topt", 128
    )
    # midori128 tenum
    midori_curve(
        0, 0, "midori_128_tenum_0_0.pdf",
        r"Solver performance ($t_{enum}$) without diff variables. Time limit: 1 hour",
        "tenum", 128
    )
    midori_curve(
        1, 0, "midori_128_tenum_1_0.pdf",
        r"Solver performance ($t_{enum}$) with diff variables. Time limit: 1 hour",
        "tenum", 128
    )