#!/bin/bash
set -u

DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
MYSQLITE3="$DIR"/sqlite3.26.0

# check number of arguments
if [ $# -ne 1 ]; then
    echo " usage: nvtxPushPopRangeStats [SQLite DB file exported from NSys QDREP file]"
    exit 1
fi

# check if file exists
if [ ! -f "$1" ]
then
    echo "$1 file not found. Exiting"
    exit 1
fi

# check if file opened is DB; if not, exit
# The sqlite3 file format is defined at https://www.sqlite.org/draft/fileformat.html
HEADER=$(head -c 16 "$1")
if [ "$HEADER" != "SQLite format 3" ]
then
    echo "$1 is not an SQlite DB file. Exiting."
    exit 1
fi

# check if DB contains nvtx PP Range data; if not, exit
COUNT=$("$MYSQLITE3" "$1" "SELECT COUNT(*) FROM NVTX_EVENTS");
if [ $COUNT -eq 0 ]
then
    echo "NVTX data was not collected."
    exit 0
fi

# check if nvtxPushPopRangeStats table exists
if [ "$("$MYSQLITE3" "$1" "SELECT name FROM sqlite_master WHERE type='table' AND name='nvtxPushPopRangeStats';")" != "nvtxPushPopRangeStats" ]
then
    printf "\nGenerating NVTX Push-Pop Range Statistics...\n"

    "$MYSQLITE3" "$1" <<GenerateSummary
    PRAGMA SYNCHRONOUS=OFF;
    -- Create a table to store domain names for easy access.
    CREATE TABLE DomainStrings (id INTEGER, domainName TEXT);
    INSERT INTO DomainStrings SELECT domainId, text from NVTX_EVENTS where eventType = 75;

    -- create a new nvtx table with only the necessary data and to avoid modifying original table
    CREATE TABLE NvtxPushPopRanges (id INTEGER, start INTEGER, end INTEGER DEFAULT 0, duration INTEGER, eventType INTEGER,
        textId INTEGER, domainId INTEGER, text TEXT, combinedTextLength INTEGER, combinedText TEXT);
    INSERT INTO NvtxPushPopRanges (id, start, eventType, textId, domainId, text)
        SELECT id, start, eventType, textId, domainId, text
        FROM NVTX_EVENTS WHERE NVTX_EVENTS.eventType = 59 OR NVTX_EVENTS.eventType = 70;

    -- Find the last nvtx range start or end - use it as the end of nvtx ranges
    --  that end after the collection ends. Assume the collection lasts less than 15 min.
    CREATE TABLE ms (maxStart INTEGER);
    INSERT INTO ms SELECT max(start) from NVTX_EVENTS where start < 900000000000;
    CREATE TABLE me (maxEnd INTEGER);
    INSERT INTO me SELECT max(end) from NVTX_EVENTS where end < 900000000000;
    CREATE TABLE estEnd (estimatedEnd INTEGER);
    INSERT INTO estEnd SELECT
        CASE
            WHEN (SELECT maxStart FROM ms) > (SELECT maxEnd FROM me) THEN (SELECT maxStart FROM ms)
            ELSE (SELECT maxEnd FROM me)
        END;

    UPDATE NvtxPushPopRanges SET end = (SELECT end from NVTX_EVENTS WHERE NVTX_EVENTS.end < 900000000000
        AND NVTX_EVENTS.id = NvtxPushPopRanges.id);
    UPDATE NvtxPushPopRanges SET end = (SELECT estimatedEnd FROM estEnd) WHERE NvtxPushPopRanges.end = 0;
    UPDATE NvtxPushPopRanges SET duration = end - start;

    -- combinedText column holds domain:range_name string where applicable.
    -- for domain combined with registered string
    UPDATE NvtxPushPopRanges SET combinedText =
        ((select domainName from DomainStrings where DomainStrings.id = NvtxPushPopRanges.domainId) ||
        ':' || (SELECT value FROM StringIds WHERE NvtxPushPopRanges.textId = StringIds.id))
        WHERE (domainId > 0 AND textId > 0) AND (eventType = 59 OR eventType = 70);
    -- for domain combined with nonregistered string
    UPDATE NvtxPushPopRanges SET combinedText = ((SELECT domainName FROM DomainStrings
        WHERE DomainStrings.id = NvtxPushPopRanges.domainId) ||
        ':' || text) WHERE (domainId > 0 AND textId = 0) AND (eventType = 59 OR eventType = 70);
    -- add the registered strings with no domain to the text column in the NvtxPushPopRanges table
    UPDATE NvtxPushPopRanges SET text =
        ((SELECT value FROM StringIds WHERE NvtxPushPopRanges.textId = StringIds.id))
        WHERE (domainId = 0 AND textId > 0) AND (eventType = 59 OR eventType = 70);


    -- Calculate the length of the combined string. Use the length to set the column width below.
    UPDATE NvtxPushPopRanges SET combinedTextLength = length(combinedText) WHERE domainId > 0;
    UPDATE NvtxPushPopRanges SET combinedTextLength = length(text) WHERE domainId = 0;

    CREATE TABLE nvtxPushPopRangeStats (tag TEXT, num INTEGER, min INTEGER, max INTEGER, avg INTEGER, total INTEGER);
    -- for non-domain case
    INSERT INTO nvtxPushPopRangeStats SELECT text, COUNT(text), min(duration), max(duration), avg(duration), sum(duration)
        FROM NvtxPushPopRanges WHERE (domainId = 0 AND (eventType = 59 OR eventType = 70)) GROUP BY text;
    -- for domain case
    INSERT INTO nvtxPushPopRangeStats SELECT combinedText, COUNT(combinedText), min(duration), max(duration), avg(duration),
        sum(duration) FROM NvtxPushPopRanges WHERE (domainId > 0 AND (eventType = 59 OR eventType = 70)) GROUP BY combinedText;
GenerateSummary
fi

printf "NVTX Push-Pop Range Statistics\n\n"
MAXSTRINGLENGTH=$("$MYSQLITE3" "$1" "SELECT max(combinedTextLength) FROM NvtxPushPopRanges");
MAXSTRINGLENGTH=$((MAXSTRINGLENGTH + 3))
MAXSTRINGLENGTH=$(( MAXSTRINGLENGTH < 100 ? MAXSTRINGLENGTH : 100 ))

TOTALTIME=$("$MYSQLITE3" "$1" "SELECT sum(total) FROM nvtxPushPopRangeStats");

echo -e ".width 0 12 0 14 12 12 $MAXSTRINGLENGTH\n SELECT round((total*100.0)/$TOTALTIME,1) as 'Time(%)',
    total as 'Time (ns)', num as Instances, round(avg,1) as 'Avg (ns)', min as 'Min (ns)', max as 'Max (ns)',
    tag as Range from nvtxPushPopRangeStats ORDER BY total DESC;" | "$MYSQLITE3" -column -header "$1"

printf "\n\n"
