#!/usr/bin/env python3

import os
import sys
import json
from math import sqrt
import numpy as np
import matplotlib.pyplot as plt
from contextlib import contextmanager

Z_STAR = 1.96

##
# This functionality allows us to temporarily change our working directory
#
# @input newdir - the new directory (relative to our current position) we want to be in
@contextmanager
def cd(newdir, makenew):
    prevdir = os.getcwd()
    directory = os.path.expanduser(newdir)
    if not os.path.exists(directory) and makenew:
        os.makedirs(directory)
    os.chdir(directory)
    try:
        yield
    finally:
        os.chdir(prevdir)

def readData(dataDirectory):
    serverData = {}
    clientData = {}

    realDirectory = os.path.expanduser(dataDirectory)
    for test in os.listdir(realDirectory):
        if not test.startswith('.') and not test.endswith('.tar.gz') and test.find("default") == -1:
            testParts = test.split("-")
            
            if not testParts[0] in serverData:
                serverData[testParts[0]] = {}

            if not testParts[0] in clientData:
                clientData[testParts[0]] = {}
            
            if not testParts[1] in serverData[testParts[0]]:
                serverData[testParts[0]][testParts[1]] = {}

            if not testParts[1] in clientData[testParts[0]]:
                clientData[testParts[0]][testParts[1]] = {}
            
            if not testParts[2] in serverData[testParts[0]][testParts[1]]:
                serverData[testParts[0]][testParts[1]][testParts[2]] = {}

            if not testParts[2] in clientData[testParts[0]][testParts[1]]:
                clientData[testParts[0]][testParts[1]][testParts[2]] = {}

            if not testParts[3] in serverData[testParts[0]][testParts[1]][testParts[2]]:
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]] = {}

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_sent'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_sent'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_sent'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_sent'] = []


            if not testParts[3] in clientData[testParts[0]][testParts[1]][testParts[2]]:
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]] = {}

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_sent'] = []

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_sent'] = []

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_sent'] = []

    for test in os.listdir(realDirectory):
        if not test.startswith('.') and not test.endswith('.tar.gz') and test.find("default") == -1:
            testParts = test.split("-")

            for whichEntity in os.listdir(os.path.join(realDirectory, test)):
                if whichEntity.startswith('s') or whichEntity.startswith('d'):
                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'overallEpoch.out')), 'r') as overallEpochFile:
                            for line in overallEpochFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'epochUp.out')), 'r') as epochUpFile:
                            for line in epochUpFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'epochDown.out')), 'r') as epochDownFile:
                            for line in epochDownFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'voteUpdate.out')), 'r') as voteUpdateFile:
                            for line in voteUpdateFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                elif whichEntity.startswith('c'):
                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'vote.out')), 'r') as repVerifierFile:
                            for line in repVerifierFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_wall'].append(wallTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_cpu'].append(cpuTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_recv'].append(dataRecv)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'repProver.out')), 'r') as repProverFile:
                            for line in repProverFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_wall'].append(wallTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_cpu'].append(cpuTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_recv'].append(dataRecv)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'repVerifier.out')), 'r') as repVerifierFile:
                            for line in repVerifierFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_wall'].append(wallTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_cpu'].append(cpuTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_recv'].append(dataRecv)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_sent'].append(dataSent)
                    except FileNotFoundError as e:
                        pass

    return serverData, clientData

def plotComparison(data, dataParts, xVariable, lineVariable, whichGraph, **kwargs):
    title = kwargs['title'] if 'title' in kwargs else ''
    xLabel = kwargs['xLabel'] if 'xLabel' in kwargs else ''
    yLabel = kwargs['yLabel'] if 'yLabel' in kwargs else ''
    fileNameStr = kwargs['fileNameStr'] if 'fileNameStr' in kwargs else f"basic-{xVariable}-{lineVariable}-{whichGraph}"
    legendLoc = kwargs['legendLoc'] if 'legendLoc' in kwargs else "best"
    legendBBoxAnchor = kwargs['legendBBoxAnchor'] if 'legendBBoxAnchor' in kwargs else (0, 1)
    loglog = kwargs['loglog'] if 'loglog' in kwargs else True
    yLim = kwargs['yLim'] if 'yLim' in kwargs else False
    aspect = kwargs['aspect'] if 'aspect' in kwargs else None

    fig = plt.figure()
    ax = fig.gca()

    whichLines = 0
    lineSelection = []
    legendStart = ""
    legendEnd = ""
    if lineVariable == 'workload':
        whichLines = 0
        lineSelection = data.keys()
        legendStart = "Workload: "
    elif lineVariable == 'numServers':
        whichLines = 1
        lineSelection = data['all'].keys()
        legendEnd = " servers"
    elif lineVariable == 'numClients':
        whichLines = 2
        lineSelection = data['all']['2'].keys()
        legendEnd = " clients"
    elif lineVariable == 'lambda':
        whichLines = 3
        lineSelection = data['all']['2']['5'].keys()
        legendStart = "Lambda: "

    whichX = 0
    xSelection = []
    if xVariable == 'workload':
        whichX = 0
        xSelection = data.keys()
    elif xVariable == 'numServers':
        whichX = 1
        xSelection = [int(x) for x in data['all'].keys()]
        xSelection.sort()
    elif xVariable == 'numClients':
        whichX = 2
        xSelection = [int(x) for x in data['all']['2'].keys()]
        xSelection.sort()
    elif xVariable == 'lambda':
        whichX = 3
        xSelection = [int(x) for x in data['all']['2']['5'].keys()]
        xSelection.sort()


    for selection in lineSelection:
        xs = []
        ys = []
        yErrs = []

        legend = legendStart + selection + legendEnd

        dataParts[whichLines] = selection
        
        for x in xSelection:
            dataParts[whichX] = str(x)

            curr_data = data[dataParts[0]][dataParts[1]][dataParts[2]][dataParts[3]][whichGraph]

            mean = np.mean(curr_data)
            std = np.std(curr_data)
            sqrt_len = sqrt(len(curr_data))

            xs.append(x)
            ys.append(mean)
            yErrs.append(Z_STAR * std / sqrt_len)

        line, _, _ = ax.errorbar(xs, ys, yerr=yErrs, capsize=7.0, label=legend)

    ax.set_title(title, fontsize='x-large')
    ax.set_xlabel(xLabel, fontsize='large')
    ax.set_ylabel(yLabel, fontsize='large')

    bottom, top = ax.get_ylim()
    bottom = (0 if bottom > 0 else bottom)
    ax.set_ylim(bottom=bottom)

    if yLim:
        ax.set_ylim(bottom=yLim[0], top=yLim[1])

    if aspect:
        ax.set_aspect(aspect, adjustable='box')

    legend = ax.legend(loc=legendLoc, bbox_to_anchor=legendBBoxAnchor, fontsize='large')

    with cd('../plt/', True):
        fig.savefig(f"{fileNameStr}.pdf", bbox_inches='tight')
    plt.close(fig)

def main(dataDirectory, plotOptionsFile):
    serverData, clientData = readData(dataDirectory)

    plotOptions = []
    with open(plotOptionsFile, 'r') as options:
        plotOptions = json.load(options)

    for option in plotOptions:
        try:
            data = serverData if (option['data'].lower() == "server" or option['data'].lower() == "s") else clientData
            dataParts = option['dataParts']
            xVariable = option['xVariable']
            lineVariable = option['lineVariable']
            whichGraph = option['whichGraph']
        except KeyError as e:
            continue
        
        kwargs = {}
        if "title" in option:
            kwargs["title"] = option["title"]
        if "xLabel" in option:
            kwargs["xLabel"] = option["xLabel"]
        if "yLabel" in option:
            kwargs["yLabel"] = option["yLabel"]
        if "fileNameStr" in option:
            kwargs["fileNameStr"] = option["fileNameStr"]
        if "legendLoc" in option:
            kwargs["legendLoc"] = option["legendLoc"]
        if "legendBBoxAnchor" in option:
            anchor = (option["legendBBoxAnchor"][0], option["legendBBoxAnchor"][1])
            kwargs["legendBBoxAnchor"] = anchor
        if "loglog" in option:
            kwargs["loglog"] = option["loglog"]
        if "yLim" in option:
            kwargs["yLim"] = option["yLim"]
        if "aspect" in option:
            kwargs["aspect"] = option["aspect"]

        plotComparison(data, dataParts, xVariable, lineVariable, whichGraph, **kwargs)

if __name__ == "__main__":
    main("../out", "../plt/plots.json")