############################ 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

### 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 --input_dir 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(profileDir, single_thread):

    infilePath = os.path.join(profileDir, 'itrace_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:
        bbl_num = int(struct.unpack('I', fIn.read(4))[0])

        #print("bbl_num = {}".format(bbl_num))

        bbls_offsets1 = list(range(bbl_num))
        bbls_offsets2 = list(range(bbl_num))

        for i in bbls_offsets1:
            bbl_id = int(struct.unpack('I', fIn.read(4))[0])
            #print("bbl_id: {:2d}".format(bbl_id))
            offset1 = int(struct.unpack('I', fIn.read(4))[0])
            bbls_offsets1[i] = offset1
            offset2 = int(struct.unpack('I', fIn.read(4))[0])
            bbls_offsets2[i] = offset2
            #print("bbls_offsets1[{:2d}] = {:8x}".format(i ,bbls_offsets1[i]))
            #print("bbls_offsets2[{:2d}] = {:8x}".format(i ,bbls_offsets2[i]))

        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("TileID = {} num_of_threads = {}".format(TileID, 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])

                #if DualSubSliceID != 0xFFFFFFFF:
                #    print('THREAD tile_{}_s_{}_dss_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, DualSubSliceID, SubSliceID, EUID, ID))
                #else:
                #    print('THREAD tile_{}_s_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, SubSliceID, EUID, ID))

                #print("num_of_records = {}".format(num_of_records))

                if DualSubSliceID != 0xFFFFFFFF:
                    outFile = os.path.join(profileDir, 'itrace___tile_{}_s_{}_dss_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, DualSubSliceID, SubSliceID, EUID, ID))
                else:
                    outFile = os.path.join(profileDir, 'itrace___tile_{}_s_{}_ss_{}_eu_{}_tid_{}.out'.format(TileID, SliceID, SubSliceID, EUID, ID))

                fOut = open(outFile,"w")

                fOut.write("  BBL ID             INS OFFSET\n"
                           "===============================\n")
                first_time = True

                while num_of_records:
                    bbl = int(struct.unpack('I', fIn.read(4))[0])
                    loopCount = int(struct.unpack('I', fIn.read(4))[0])
                    #print("bbl: {:2d} loopCount = {}".format(bbl, loopCount))
                    offset1 = bbls_offsets1[bbl]
                    offset2 = bbls_offsets2[bbl]
                    if not first_time and bbl == 0:
                        fOut.write("--- EOT ---\n")

                    fOut.write("  {:4d}                {:#04x}\n".format(bbl, offset1));

                    if offset1 != offset2 and offset2 != 0xFFFFFFFF:
                        if loopCount > 1:
                            fOut.write("  {:4d}                {:#04x}  -   repeats {} times\n".format(bbl, offset2, loopCount));
                        else:
                            fOut.write("  {:4d}                {:#04x}\n".format(bbl, offset2));

                    num_of_records = num_of_records - 1

                    if first_time:
                        first_time = False

                fOut.write("--- EOT ---\n")

                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("--extract-one-file",  dest="extract_one_file",default=False,
                        action='store_true',
                        help="forcing extraction single file from trace")

    args = parser.parse_args()

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

    ProcessProfileDir(args.input_dir, args.extract_one_file)

Main()
