############################ begin_copyright_notice ############################
### Copyright (C) 2018-2022 Intel Corporation
### 
### SPDX-License-Identifier: MIT
############################ end_copyright_notice ##############################

#!/usr/bin/python

import sys, re, tempfile, os, os.path, shutil, time, struct
import array

from array import *
from argparse import ArgumentParser

TRACE_TYPE_DICT = {
    "funtime"    : 0,
    "memlatency" : 1,
    "occupancy"  : 2
}

TRACE_TYPE_TO_STRING = {
    0 : "funtime",
    1 : "memlatency",
    2 : "occupancy"
}

### My print functions
def Print(msg):
    print('%s' %msg)
    sys.__stdout__.flush()

def Info(msg):
    Print(msg)

def Prompt(msg):
    print('%s' %msg)
    sys.__stdout__.flush()

def Error(msg):
    MyExit('Error: %s' %msg)

### Assertions
def AssertAlways(cond, msg):
    if not cond:
        Print('Assert failure: %s' %msg)
        sys.exit()

def AssertOpt(cond, msg):
    AssertAlways(cond, msg)

def MyExit(msg):
    Print(msg)
    sys.exit()

def Usage(progName):
    return '\nUsage: %s [--occupancy|--funtime|--memlatency] c:\profile_dir\Session_Final\PS_asm5782fb10f48402ef_simd8\draw__1359' %(progName)

def ParseArgs():
    progName = sys.argv[0]

    if (len(sys.argv) != 2):
        Error(Usage(progName))

    profilerDir = sys.argv[1]

    return (profilerDir)

        
def ProcessProfileDir(input_dir, trace_type, single_thread):
    if not os.path.isdir(input_dir):
        Error('Directory "%s" does not exist'%input_dir)
        
    infilePath = os.path.join(input_dir, 'perftrace_compressed.bin')
   
    if not os.path.isfile(infilePath):
        Error('File "%s" does not exist'%infilePath)


    try:
        fIn = open(infilePath, 'rb')
    except (IOError, e):
        Error(e)
    else:
        
        trace_mode  = int(struct.unpack('<I', fIn.read(4))[0])
        
        if TRACE_TYPE_DICT[trace_type] != trace_mode:
            print("The trace file perftrace_compressed.bin is of a wrong type '{}' (configured type '{}')".format(TRACE_TYPE_TO_STRING[trace_mode], trace_type))
        
        num_of_tiles = int(struct.unpack('I', fIn.read(4))[0])
        #print("num_of_tiles = {}".format(num_of_tiles))

        while num_of_tiles:

            TileID = int(struct.unpack('I', fIn.read(4))[0])

            num_of_threads = int(struct.unpack('I', fIn.read(4))[0])      
            #print("num_of_threads = {}".format(num_of_threads))
        
            while num_of_threads:
                SliceID = int(struct.unpack('I', fIn.read(4))[0])
                DualSubSliceID = int(struct.unpack('I', fIn.read(4))[0])
                SubSliceID = int(struct.unpack('I', fIn.read(4))[0])
                EUID = int(struct.unpack('I', fIn.read(4))[0])
                ID = int(struct.unpack('I', fIn.read(4))[0])
                num_of_records = int(struct.unpack('I', fIn.read(4))[0])
            
                #print("tid = {}".format(tid))
                #print("num_of_records = {}".format(num_of_records))

                #outFile = os.path.join(input_dir, 'itrace_{}.out'.format(tid))
                if DualSubSliceID != 0xFFFFFFFF:
                    outFile = os.path.join(input_dir, '{}___tile_{}_s_{}_dss_{}_ss_{}_eu_{}_tid_{}.out'.format(trace_type, TileID, SliceID, DualSubSliceID, SubSliceID, EUID, ID))
                else:
                    outFile = os.path.join(input_dir, '{}___tile_{}_s_{}_ss_{}_eu_{}_tid_{}.out'.format(trace_type, TileID, SliceID, SubSliceID, EUID, ID))
                fOut = open(outFile,"w")

                if trace_type == "occupancy":
                    fOut.write(" Invocation   Start          End\n"
                               "================================\n")
                else:
                    fOut.write(" Invocation   Duration in cycles\n"
                               "================================\n")
                s = 0
            
                while num_of_records:
            
                    duration = 0
                
                    if trace_type == "memlatency":
                        duration = int(struct.unpack('<I', fIn.read(4))[0])
                    else:
                        before = int(struct.unpack('<Q', fIn.read(8))[0])
                        after  = int(struct.unpack('<Q', fIn.read(8))[0])
                
                        if before < after:
                            duration = after - before
                        else:
                            duration = 0xFFFFFFFF


                    str = ""

                    if trace_type == "occupancy":
                        str = "  {:6d}        {:#016x}      {:#016x}         duration = {:d}".format(s,before, after, duration)
                    elif trace_type == "funtime":
                        str = "  {:6d}        {:d}".format(s, duration)
                    else:
                        if duration == 0xFFFFFFFF:
                            str = "  {:6d}        Skip".format(s)
                        else:
                            str = "  {:6d}        {:d}".format(s,duration)

                    fOut.write("{}\n".format(str))
                                        
                    num_of_records = num_of_records - 1
                    s = s + 1

                fOut.close()

                if single_thread:
                    break

                num_of_threads = num_of_threads - 1

            num_of_tiles = num_of_tiles - 1

        fIn.close()

def Main():
    
    
    parser = ArgumentParser()
    parser.add_argument("--input_dir",  dest="input_dir",default='',
                        help="write report to FILE", metavar="FILE")
    parser.add_argument("--occupancy", 
                        action="store_true", dest="occupancy", default=False,
                        help="uncompress occupancy trace")
    parser.add_argument("--memlatency", 
                        action="store_true", dest="memlatency", default=False,
                        help="uncompress memory latency trace")
    parser.add_argument("--funtime", 
                        action="store_true", dest="funtime", default=False,
                        help="uncompress funtime trace")
    parser.add_argument("--extract-one-file",  dest="extract_one_file",default=False,
                        action='store_true',
                        help="forcing extraction single file from trace")
    args = parser.parse_args()
       
    print('input_dir  = {}'.format(args.input_dir))
    print('occupancy  = {}'.format(args.occupancy))
    print('memlatency = {}'.format(args.memlatency))
    print('funtime    = {}'.format(args.funtime))

    if (not os.path.exists(args.input_dir)):
        print("the directory {} doesn't exist".format(args.input_dir))
        Error(Usage(sys.argv[0]))

    if not args.occupancy and not args.funtime and not args.memlatency:
        print("the trace type should be provided")
        Error(Usage(sys.argv[0]))

    if args.occupancy and (args.funtime or args.memlatency):
        print("single trace type should be provided")
        Error(Usage(sys.argv[0]))

    if args.funtime and args.memlatency:
        print("single trace type should be provided")
        Error(Usage(sys.argv[0]))

    if args.occupancy:
        trace_type = 'occupancy'

    if args.funtime:
        trace_type = 'funtime'

    if args.memlatency:
        trace_type = 'memlatency'

    ProcessProfileDir(args.input_dir, trace_type, args.extract_one_file)
    
Main()
