import sys
import sqlite3
import ntpath
import os
import matplotlib.pyplot as plt
import numpy as np
from memUtil import *
from math import *
from matplotlib.backends.backend_pdf import PdfPages

numberOfBins = "auto"
latencyRange = None

plots = []

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

def getThreads(connection):
    cursor = connection.cursor()
    cursor.execute("SELECT DISTINCT(Thread) FROM transactions WHERE Thread != 0 ORDER BY Thread")
    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:

    # The data transfer begins and ends inside the time window
    queryFull = """
        SELECT
            SUM(DataStrobeEnd - DataStrobeBegin)
        FROM
            transactions
        WHERE
            DataStrobeBegin >= ?
            AND DataStrobeEnd <= ?
            AND TThread = {0}
    """

    # Only the end of the data transfer is inside the time window
    queryEnd = """
        SELECT
            SUM(DataStrobeEnd - ? )
        FROM
            transactions
        WHERE
            DataStrobeBegin < ?
            AND DataStrobeEnd > ?
            AND DataStrobeEnd <=?
            AND TThread = {0}
    """

    # Only the beginning of the data transfer is inside the time window
    queryBegin = """
        SELECT
            SUM( ? - DataStrobeBegin)
        FROM
            transactions
        WHERE
            DataStrobeBegin >= ?
            AND DataStrobeBegin < ?
            AND DataStrobeEnd > ?
            AND TThread = {0}
    """

    # The data transfer occupies all the time window
    queryPart = """
        SELECT
            DataStrobeBegin
        FROM
            transactions
        WHERE
            DataStrobeBegin <= ?
            AND DataStrobeEnd >= ?
            AND TThread = {0}
    """

    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

@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:

    # The data transfer begins and ends inside the time window
    queryFull = """
        SELECT
            SUM(DataStrobeEnd - DataStrobeBegin)
        FROM
            Phases
        WHERE
            DataStrobeBegin >= ?
            AND DataStrobeEnd <= ?
    """

    # Only the end of the data transfer is inside the time window
    queryEnd = """
        SELECT
            SUM(DataStrobeEnd - ?)
        FROM
            Phases
        WHERE
            DataStrobeBegin < ?
            AND DataStrobeEnd > ?
            AND DataStrobeEnd <= ?
    """

    # Only the beginning of the data transfer is inside the time window
    queryBegin = """
        SELECT
            SUM(? - DataStrobeBegin)
        FROM
            Phases
        WHERE
            DataStrobeBegin >= ?
            AND DataStrobeBegin < ?
            AND DataStrobeEnd > ?
    """

    # The data transfer occupies all the time window
    queryPart = """
        SELECT
            DataStrobeBegin
        FROM
            Phases
        WHERE
            DataStrobeBegin <= ?
            AND DataStrobeEnd >= ?
    """

    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)

    # 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 response_latency_window(connection, tracePath, steps):
    windowSize = getWindowSize(connection)

    cursor = connection.cursor()
    query = """
        SELECT
        avg(RESP.PHASEBEGIN - REQ.PHASEBEGIN) / 1000
        FROM
        PHASES REQ,
        PHASES RESP
        WHERE
        REQ.PHASENAME = 'REQ'
        AND RESP.PHASENAME = 'RESP'
        AND REQ.TRANSACT = RESP.TRANSACT
        AND RESP.PHASEBEGIN >= ? and RESP.PHASEEND <= ?
    """
    outputFileName, basename = createOutputFilename(tracePath, 'response_latency', '', 'pdf')
    outputFile = "{0}\n\t".format(outputFileName)

    LatencyFigure = plt.figure(figsize=(10, 5), dpi=300)
    LatencyFigurePlot = LatencyFigure.add_subplot(111)
    LatencyFigurePlot.set_xlabel('Time [ns]')
    LatencyFigurePlot.set_ylabel('Response Latency [ns]')
    LatencyFigurePlot.set_title('Average Response Latency: ' + str(basename))
    LatencyFigurePlot.grid(True)

    time = [None] * steps
    latency = [None] * steps

    for i in range(steps):
        cursor.execute(query, (i * windowSize, (i + 1) * windowSize))
        result = cursor.fetchone()[0]
        time[i] = ((i * windowSize) - (windowSize / 2)) / 1000 # ps to ns
        latency[i] = result

    LatencyFigurePlot.plot(time, latency, linewidth=0.5, label="Latency")
    LatencyFigurePlot.legend(loc="upper left")

    pdf = PdfPages(outputFileName)
    pdf.savefig(LatencyFigure)
    pdf.close()
    LatencyFigurePlot.clear()
    plt.close()

    return outputFile

@plot
def wr_response_latency_window(connection, tracePath, steps):
    windowSize = getWindowSize(connection)

    cursor = connection.cursor()
    query = """
        SELECT
        avg(RESP.PHASEBEGIN - REQ.PHASEBEGIN) / 1000
        FROM
        PHASES REQ,
        PHASES RESP
        INNER JOIN
            Transactions
            ON REQ.TRANSACT = Transactions.ID
        WHERE
        REQ.PHASENAME = 'REQ'
        AND RESP.PHASENAME = 'RESP'
        AND REQ.TRANSACT = RESP.TRANSACT
        AND RESP.PHASEBEGIN >= ? and RESP.PHASEEND <= ?
        AND Transactions.Command = "W"
    """
    outputFileName, basename = createOutputFilename(tracePath, 'wr_response_latency', '', 'pdf')
    outputFile = "{0}\n\t".format(outputFileName)

    LatencyFigure = plt.figure(figsize=(10, 5), dpi=300)
    LatencyFigurePlot = LatencyFigure.add_subplot(111)
    LatencyFigurePlot.set_xlabel('Time [ns]')
    LatencyFigurePlot.set_ylabel('Response Latency [ns]')
    LatencyFigurePlot.set_title('Average Write Response Latency: ' + str(basename))
    LatencyFigurePlot.grid(True)

    time = [None] * steps
    latency = [None] * steps

    for i in range(steps):
        cursor.execute(query, (i * windowSize, (i + 1) * windowSize))
        result = cursor.fetchone()[0]
        time[i] = ((i * windowSize) - (windowSize / 2)) / 1000 # ps to ns
        latency[i] = result

    LatencyFigurePlot.plot(time, latency, linewidth=0.5, label="Latency")
    LatencyFigurePlot.legend(loc="upper left")

    pdf = PdfPages(outputFileName)
    pdf.savefig(LatencyFigure)
    pdf.close()
    LatencyFigurePlot.clear()
    plt.close()

    return outputFile

@plot
def rd_response_latency_window(connection, tracePath, steps):
    windowSize = getWindowSize(connection)

    cursor = connection.cursor()
    query = """
        SELECT
        avg(RESP.PHASEBEGIN - REQ.PHASEBEGIN) / 1000
        FROM
        PHASES REQ,
        PHASES RESP
        INNER JOIN
            Transactions
            ON REQ.TRANSACT = Transactions.ID
        WHERE
        REQ.PHASENAME = 'REQ'
        AND RESP.PHASENAME = 'RESP'
        AND REQ.TRANSACT = RESP.TRANSACT
        AND RESP.PHASEBEGIN >= ? and RESP.PHASEEND <= ?
        AND Transactions.Command = "R"
    """
    outputFileName, basename = createOutputFilename(tracePath, 'rd_response_latency', '', 'pdf')
    outputFile = "{0}\n\t".format(outputFileName)

    LatencyFigure = plt.figure(figsize=(10, 5), dpi=300)
    LatencyFigurePlot = LatencyFigure.add_subplot(111)
    LatencyFigurePlot.set_xlabel('Time [ns]')
    LatencyFigurePlot.set_ylabel('Response Latency [ns]')
    LatencyFigurePlot.set_title('Average Read Response Latency: ' + str(basename))
    LatencyFigurePlot.grid(True)

    time = [None] * steps
    latency = [None] * steps

    for i in range(steps):
        cursor.execute(query, (i * windowSize, (i + 1) * windowSize))
        result = cursor.fetchone()[0]
        time[i] = ((i * windowSize) - (windowSize / 2)) / 1000 # ps to ns
        latency[i] = result

    LatencyFigurePlot.plot(time, latency, linewidth=0.5, label="Latency")
    LatencyFigurePlot.legend(loc="upper left")

    pdf = PdfPages(outputFileName)
    pdf.savefig(LatencyFigure)
    pdf.close()
    LatencyFigurePlot.clear()
    plt.close()

    return outputFile

@plot
def command_bus_utilisation_window(connection, tracePath, steps):
    windowSize = getWindowSize(connection)

    cursor = connection.cursor()

    # Query that sums all lengths that are completely contained in the window
    query_full = """
        SELECT
            SUM(CommandLengths.Length)
        FROM
            Phases,
            GeneralInfo
            INNER JOIN
                CommandLengths
                ON Phases.PhaseName = CommandLengths.Command
            WHERE
                Phases.PhaseBegin >= ?
                AND (Phases.PhaseBegin + (CommandLengths.Length * GeneralInfo.clk)) < ?
    """

    # Gets the PhaseBegin of the command that reaches out of the window
    # query_border = """
    #     SELECT
    #         Phases.PhaseBegin
    #     FROM
    #         Phases,
    #         GeneralInfo
    #         INNER JOIN
    #             CommandLengths
    #             ON Phases.PhaseName = CommandLengths.Command
    #         WHERE
    #             Phases.PhaseBegin >= ?
    #             AND Phases.PhaseBegin < ?
    #             AND (Phases.PhaseBegin + (CommandLengths.Length * GeneralInfo.clk)) >= ?
    # """

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

    LatencyFigure = plt.figure(figsize=(10, 5), dpi=300)
    LatencyFigurePlot = LatencyFigure.add_subplot(111)
    LatencyFigurePlot.set_xlabel('Time [ns]')
    LatencyFigurePlot.set_ylabel('Utilization [%]')
    LatencyFigurePlot.set_title('Command Bus Utilization: ' + str(basename))
    LatencyFigurePlot.grid(True)

    clk, _ = getClock(connection)

    time = [None] * steps
    utilization = [None] * steps

    for i in range(steps):
        left_limit = i * windowSize
        right_limit = (i + 1) * windowSize

        cursor.execute(query_full, (left_limit, right_limit))
        result = cursor.fetchone()[0]

        if (result is None):
            result = 0

        cmdBusOccupied = result * clk

        time[i] = ((i * windowSize) - (windowSize / 2)) / 1000 # ps to ns
        utilization[i] = cmdBusOccupied / windowSize * 100

        if (utilization[i] > 100):
            print(left_limit, right_limit)

    LatencyFigurePlot.plot(time, utilization, linewidth=0.5, label="Utilization")
    LatencyFigurePlot.legend(loc="upper left")

    pdf = PdfPages(outputFileName)
    pdf.savefig(LatencyFigure)
    pdf.close()
    LatencyFigurePlot.clear()
    plt.close()

    return outputFile

@plot
def queue_window(connection, tracePath, steps):
    cursor = connection.cursor()
    cursor.execute("select max(BufferNumber) from BufferDepth;")
    bufferNumber = int(cursor.fetchone()[0]) + 1

    cursor = connection.cursor()
    cursor.execute("select MaxBufferDepth from GeneralInfo;")
    maxBufferDepth = int(cursor.fetchone()[0])

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

    QueueFigure = plt.figure(figsize=(10, 5), dpi=300)
    QueueFigurePlot = QueueFigure.add_subplot(111)
    QueueFigurePlot.set_xlabel('Time [s]')
    QueueFigurePlot.set_ylabel('Queue Utilization')
    QueueFigurePlot.set_title('Average Queue Utilization: ' + str(basename))
    QueueFigurePlot.grid(True)


    for b in range(bufferNumber):
        cursor.execute("select Time, AverageBufferDepth from BufferDepth where BufferNumber = {};".format(b))
        time = [None] * steps
        queue = [None] * steps
        for i in range(steps-1):
            result = cursor.fetchone()
            time[i] = result[0]
            queue[i] = result[1]

        QueueFigurePlot.plot(time, queue, linewidth=0.5, label="Queue {}".format(b))

    QueueFigurePlot.legend(loc="upper left")

    x1,x2,y1,y2 = QueueFigurePlot.axis()
    QueueFigurePlot.axis((x1,x2,0,maxBufferDepth))

    pdf = PdfPages(outputFileName)
    pdf.savefig(QueueFigure)
    pdf.close()
    QueueFigurePlot.clear()
    plt.close()

    return outputFile

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

    windowSize = getWindowSize(connection)

    outputFile = ""
    cursor = connection.cursor()

    cursor.execute(" SELECT * FROM Power")

    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)

    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

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)