import sys
import sqlite3
from memUtil import *
from math import *
import ntpath
import os
numberOfBins = "auto"
latencyRange = None

plots = []


def getThreads(connection):
    cursor = connection.cursor()
    cursor.execute("SELECT DISTINCT(TThread) FROM transactions WHERE TThread != 0 ORDER BY TThread")
    result = []
    for currentRow in cursor:
        result.append(currentRow[0])
    return result


def createOutputFilename(tracePath, plot_type, target_measurement, file_type):
    name = ntpath.basename(tracePath)
    basename, extension = os.path.splitext(name)
    outputFileName = plot_type + '_' + target_measurement + basename + '.' + file_type
    return outputFileName, target_measurement + basename


def accessDatabase(connection, query):
    cursor = connection.cursor()
    # cursor.execute(" ")
    cursor.execute(query)
    resultArray = []
    while True:
        result = cursor.fetchone()
        if (result is not None):
            resultArray.append(result[0])
        else:
            break
    return resultArray


def calculate_bandwidth_util(connection, windowSize, steps, queryFull, queryEnd, queryBegin, queryPart):
    cursor = connection.cursor()
    maxDataRate = maximum_data_rate(connection)
    maximumPercentage = 0
    bandwidthPercentage = [0] * (steps+1)
    bandwidth = [0] * (steps+1)
    bandwidthPercentage[0] = 0
    bandwidth[0] = 0
    for i in range(steps):
        bandwidthPercentage[i+1] = 0
        cursor.execute(queryPart, (i*windowSize, (i+1)*windowSize))
        result = cursor.fetchone()
        if(result is None):
            cursor.execute(queryFull, (i*windowSize, (i+1)*windowSize))
            result = cursor.fetchone()
            if(result[0] is not None):
                bandwidthPercentage[i+1] += int(result[0])
            cursor.execute(queryEnd, (i*windowSize, i*windowSize, i*windowSize, (i+1)*windowSize))
            result = cursor.fetchone()
            if(result[0] is not None):
                bandwidthPercentage[i+1] += int(result[0])
            cursor.execute(queryBegin, ((i+1)*windowSize, i*windowSize, (i+1)*windowSize, (i+1)*windowSize))
            result = cursor.fetchone()
            if(result[0] is not None):
                bandwidthPercentage[i+1] += int(result[0])
        else:
            bandwidthPercentage[i+1] = windowSize
        bandwidthPercentage[i+1] = float(bandwidthPercentage[i+1]/windowSize)
        bandwidth[i+1] = float(bandwidthPercentage[i+1])*float(maxDataRate)/1024
        bandwidthPercentage[i+1] *= 100
        if(maximumPercentage < 100 and maximumPercentage < bandwidthPercentage[i+1]):
            maximumPercentage = bandwidthPercentage[i+1]

    return bandwidthPercentage, bandwidth, maximumPercentage


def memory_utilisation_window_thread(connection, tracePath, steps, thread_ID):
    # All possible cases of data transfers inside a time window
    queryFull = """ SELECT sum(DataStrobeEnd - DataStrobeBegin) FROM transactions Where DataStrobeBegin >= ? and DataStrobeEnd <= ? and TThread = {0}"""  # The data transfer begins and ends inside the time window
    queryEnd = """ SELECT sum(DataStrobeEnd - ?) FROM transactions Where DataStrobeBegin < ? and  DataStrobeEnd > ? and DataStrobeEnd <=? and TThread = {0}"""  # Only the end of the data transfer is inside the time window
    queryBegin = """ SELECT sum(? - DataStrobeBegin) FROM transactions Where DataStrobeBegin >= ? and DataStrobeBegin <  ? and DataStrobeEnd > ? and TThread = {0}"""  # Only the beginning of the data transfer is inside the time window
    queryPart = """ SELECT DataStrobeBegin FROM transactions Where DataStrobeBegin <= ? and DataStrobeEnd >= ? and TThread = {0}"""  # The data transfer occupies all the time window

    queryFull = queryFull.format(thread_ID)
    queryEnd = queryEnd.format(thread_ID)
    queryBegin = queryBegin.format(thread_ID)
    queryPart = queryPart.format(thread_ID)

    windowSize = getWindowSize(connection)

    bandwidthPercentage, bandwidth, maximumPercentage = calculate_bandwidth_util(connection, windowSize, steps, queryFull, queryEnd, queryBegin, queryPart)

    outputFileNameBWMatlab, basename = createOutputFilename(tracePath, 'memory_utilization_percent', 'thread_' + str(thread_ID) + '_', 'txt')
    return bandwidthPercentage, bandwidth, outputFileNameBWMatlab


def plot(function):
    plots.append(function)
    return function


@plot
def memory_utilisation_window(connection, tracePath, steps):
    # This function determines the average memory bandwidth over time in
    # percentage and in Gbit/s. The average bandwidth over time is done
    # dividing the time into windows of the same length and getting the average
    # bandwidth in each window. Through data from the database, DataStrobeEnd
    # and DataStrobeBegin, it is possible to access when a data transfer begins
    # or ends. Hence, it is achievable to check when a data transfer happens
    # and if it occupies or is inside a time window. Then, it is attainable to
    # determine the average bandwidth in percentage. Besides, extracting the
    # data from the memory specs, it is feasible to calculate the maximum data
    # rate of the memory and then determine the bandwidth in Gbit/s. The
    # bandwidth data are then plotted in two graphics.

    windowSize = getWindowSize(connection)
    maxDataRate = maximum_data_rate(connection)

    # All possible cases of data transfers inside a time window
    queryFull = """ SELECT sum(DataStrobeEnd - DataStrobeBegin) FROM transactions Where DataStrobeBegin >= ? and DataStrobeEnd <= ?"""  # The data transfer begins and ends inside the time window
    queryEnd = """ SELECT sum(DataStrobeEnd - ?) FROM transactions Where DataStrobeBegin < ? and  DataStrobeEnd > ? and DataStrobeEnd <=?"""  # Only the end of the data transfer is inside the time window
    queryBegin = """ SELECT sum(? - DataStrobeBegin) FROM transactions Where DataStrobeBegin >= ? and DataStrobeBegin <  ? and DataStrobeEnd > ?"""  # Only the beginning of the data transfer is inside the time window
    queryPart = """ SELECT DataStrobeBegin FROM transactions Where DataStrobeBegin <= ? and DataStrobeEnd >= ?"""  # The data transfer occupies all the time window

    bandwidthPercentage, bandwidth, maximumPercentage = calculate_bandwidth_util(connection, windowSize, steps, queryFull, queryEnd, queryBegin, queryPart)

    outputFileNameGBPS, basename = createOutputFilename(tracePath, 'memory_utilization_gbps', '', 'pdf')
    outputFileNamePercent, basename = createOutputFilename(tracePath, 'memory_utilization_percent', '', 'pdf')
    outputFileNameBWMatlab, basename = createOutputFilename(tracePath, 'memory_utilization_percent', '', 'txt')
    outputFiles = "{0}\n\t{1}\n\t{2}\n\t".format(outputFileNameGBPS, outputFileNamePercent, outputFileNameBWMatlab)

    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib.backends.backend_pdf import PdfPages

    # windowSize/1000: picoseconds to nanoseconds conversion
    time = np.arange(0, (steps+1)*windowSize/1000, windowSize/1000)
    maxBandwidth = [maxDataRate/1024] * (steps+1)

    f = open(outputFileNameBWMatlab, 'w')
    for i in range(steps):
        line = "{} {}\n".format(time[i], bandwidthPercentage[i])
        f.write(line)

    # Plot Bandwidth in Percent
    BWPercentageFigure = plt.figure()
    BWPercentageFigurePlot = BWPercentageFigure.add_subplot(111)
    BWPercentageFigurePlot.set_xlabel('Time [ns]')
    BWPercentageFigurePlot.set_ylabel('Bandwidth [%]')
    BWPercentageFigurePlot.set_ylim(-1, maximumPercentage + (10 - maximumPercentage % 10))
    BWPercentageFigurePlot.set_title('Memory Utilization in % ' + str(basename))
    BWPercentageFigurePlot.grid(True)
    BWPercentageFigurePlot.plot(time, bandwidthPercentage, label='Total')
    BWPercentageFigurePlot.legend(loc="upper left")

    # Plot absolute bandwidth
    BWFigure = plt.figure()
    BWFigurePlot = BWFigure.add_subplot(111)
    BWFigurePlot.set_xlabel('Time [ns]')
    BWFigurePlot.set_ylabel('Bandwidth [Gibit/s]')
    BWFigurePlot.set_title('Memory Utilization in Gbps ' + str(basename))
    BWFigurePlot.grid(True)
    BWFigurePlot.plot(time, bandwidth, label='Total')
    BWFigurePlot.legend(loc="upper left")

    # plt.ylim((-0.01)*float(maxDataRate)/1024, ((maximumPercentage + (10 - maximumPercentage%10))/100)*float(maxDataRate)/1024)

    threads = getThreads(connection)
    if (len(threads) > 1):
        for thread in threads:
            threadStr = "Thread " + str(thread)
            bandwidthPercentage, bandwidth, outputFileNameBWMatlab = memory_utilisation_window_thread(connection, tracePath, steps, thread)
            BWPercentageFigurePlot.plot(time, bandwidthPercentage, label=threadStr)
            BWPercentageFigurePlot.legend(loc="upper left")
            BWFigurePlot.plot(time, bandwidth, label=threadStr)
            BWFigurePlot.legend(loc="upper left")
            f = open(outputFileNameBWMatlab, 'w')
            for i in range(steps):
                line = "{} {}\n".format(time[i], bandwidthPercentage[i])
                f.write(line)
            outputFiles += "{0}\n\t".format(outputFileNameBWMatlab)

    # Save to PDF files
    pdf = PdfPages(outputFileNamePercent)
    pdf.savefig(BWPercentageFigure)
    pdf.close()
    BWPercentageFigure.clear()
    pdf = PdfPages(outputFileNameGBPS)
    BWFigurePlot.plot(time, maxBandwidth)
    pdf.savefig(BWFigure)
    pdf.close()
    BWFigurePlot.clear()
    plt.close()

    return outputFiles


@plot
def power_window(connection, tracePath, steps):

    windowSize = getWindowSize(connection)

    outputFile = ""
    cursor = connection.cursor()

    cursor.execute(" SELECT * FROM Power")

    import numpy as np
    window = float(windowSize) / pow(10, 12)
    time = np.arange(0, (windowSize * (steps + 1)) / pow(10, 6), windowSize / pow(10, 6))
    power = np.full(len(time), 0)

    for i in range(steps):
        sum = 0.0
        counter = 0
        result = cursor.fetchone()

        while (result is not None):
            sum += float(result[1])
            counter = counter + 1
            if(result[0] > window*i):
                break
            result = cursor.fetchone()

        if(counter == 0):
            break

        sum = sum / counter
        power[i+1] = sum

    outputFileName, basename = createOutputFilename(tracePath, 'power', '', 'pdf')
    outputFile = "{0}\n\t".format(outputFileName)

    import matplotlib.pyplot as plt
    from matplotlib.backends.backend_pdf import PdfPages

    PowFigure = plt.figure(figsize=(10, 5), dpi=300)
    PowFigurePlot = PowFigure.add_subplot(111)
    PowFigurePlot.set_xlabel('Time [us]')
    PowFigurePlot.set_ylabel('Power [mW]')
    PowFigurePlot.set_title('Power Consumption ' + str(basename))
    PowFigurePlot.grid(True)
    PowFigurePlot.plot(time, power, linewidth=0.5)

    pdf = PdfPages(outputFileName)
    pdf.savefig(PowFigure)
    pdf.close()
    PowFigurePlot.clear()
    plt.close()

    return outputFile


@plot
def latency_histogram(connection, tracePath, steps):
    # This function plots an histogram with access latencys

    def plot_latency_histogram(dataArray, outputFileName, basename):
        # plot into PDF file
        import matplotlib
        matplotlib.use('TkAgg')
        import matplotlib.pyplot as plt
        from matplotlib.backends.backend_pdf import PdfPages

        plt.hist(dataArray, bins=numberOfBins, range=latencyRange, histtype='bar', facecolor='green')
        plt.grid(True)
        plt.xlabel("Access Time [ns]")
        plt.ylabel("Number of Accesses (Frequency)")
        plt.title("Latency Histogram " + str(basename))
        pdf = PdfPages(outputFileName)
        pdf.savefig()
        pdf.close()
        plt.close()

    def create_latency_hist(connection, tracePath, target_measurement, query):
        # form output file name
        ofbname = 'access_latency_hist'
        outputFileName, basename = createOutputFilename(tracePath, ofbname, target_measurement, 'pdf')
        # return log string
        outputFile = "{0}\n\t".format(outputFileName)
        # access database
        resultArray = accessDatabase(connection, query)
        # plot
        plot_latency_histogram(resultArray, outputFileName, basename)
        return outputFile

    # create overal latency histogram
    query = """ SELECT ((p2.PhaseEnd - p1.PhaseBegin)/1000)
                FROM Transactions t, Phases p1, Phases p2
                WHERE t.id = p1.Transact
                    AND t.id = p2.Transact
                    AND p1.PhaseName = "REQ"
                    AND p2.PhaseName = "RESP" """
    outputFile = create_latency_hist(connection, tracePath, '', query)

    # create per-thread latency histogram
    threads = getThreads(connection)
    if (len(threads) > 1):
        queryThread = """   SELECT ((p2.PhaseEnd - p1.PhaseBegin)/1000)
                            FROM Transactions t, Phases p1, Phases p2
                            WHERE t.id = p1.Transact
                                AND t.id = p2.Transact
                                AND p1.PhaseName = "REQ"
                                AND p2.PhaseName = "RESP"
                                AND t.TThread = {0} """
        for thread in threads:
            thrname = 'thread_' + str(thread) + '_'
            outputFile += create_latency_hist(connection, tracePath, thrname, queryThread.format(thread))

    return outputFile


def generatePlots(pathToTrace):
    connection = sqlite3.connect(pathToTrace)

    print("================================")
    print("Generating plots for {0}".format(pathToTrace))

    outputFiles = "Output files are:\n\t"
    cursor = connection.cursor()
    cursor.execute(" SELECT WindowSize FROM GeneralInfo")
    windowSize = float(cursor.fetchone()[0])
    if(windowSize == 0):
        outputFiles = "No output file created. Check WindowSize and EnableWindowing configs."
    else:
        cursor.execute(" SELECT TraceEnd FROM GeneralInfo ")
        traceEnd = float(cursor.fetchone()[0])
        steps = int(ceil(traceEnd/windowSize))
        for p in plots:
            outputFiles += p(connection, pathToTrace, steps)

    connection.close()

    print(outputFiles)

    return outputFiles


if __name__ == "__main__":
    path = sys.argv[1]
    if ((len(sys.argv)) > 2):
        latencyRange = (0, int(sys.argv[2]))  # Optional argument to use a different range
        if ((len(sys.argv)) > 3):
            numberOfBins = int(sys.argv[3])  # Optional argument to use a different number of bins
    generatePlots(path)
