#!/usr/bin/python3
# Make sure to run this as root!!!!!
import os
import sys
import re
import csv
import time
import subprocess

A_SUITE = 1
A_BIN = 2
A_CRIT = 3
A_PERIOD = 5

P_CORE = 0
P_CRIT = 1
P_T1_ID = 2
P_T2_ID = 5

T_ID = 0
T_BIN = 1
T_PERIOD = 2
T_CCX = 3

CRIT_LEVEL_C = 2

all_pids=[]
def run(command):
    print(command)
    os.system(command)

def addpid(pid):
    with open("./pids.txt", "a") as f:
        f.write(str(pid) + "\n")
def ID2PID(lookupID):
    res = None
    # Find benchmark PID (avoid getting shell or numactl PID)
    try:
        res = subprocess.check_output("pgrep -af ' " + lookupID + " ' | grep -v numactl", shell=True)
    except:
        res = None
    if res:
        # Ignore the newline
        return res.decode("utf-8").split(" ")[0].strip()
    else:
        return None

def main(pathName):
    with open("./pids.txt", "w") as f:
        f.write("")

    tacle_pairs_path = "./all_pairs"
    tacle_baseline_path = "./baseline"
    dis_path = "./dis"
    # List of lists
    # Each list represents the parameters for a task as:
    # [id, suite, binary name, criticality level (0-base), pet us (unused), period ms, wss (unused)]
    all_tasks = []
    # List of lists
    # Each list represents the parameters for a pairing as:
    # [core, criticality level (0-base), task 1 ID, task 1 name (unused), task 1 period, task 2 ID, task 2 name, task 2 period]
    # if task 1 ID = task 2 ID, it's not actually a pair at all
    levelAB = []
    levelC = []
    levelC_s = []
    input_cmd = {}
    # Load input command for each binary
    with open("baseline/tacleNames.txt") as f:
        for benchName in f:
            # TACLeBench doesn't need inputs
            input_cmd[benchName.strip()] = "echo ' '"
    with open("SD-VBS/sd-vbsNames.txt") as f:
        for line in f:
            name, cmd = line.split(maxsplit=1)
            input_cmd[name.replace("./", "")] = cmd.replace("./", "./SD-VBS/").strip()
    with open("dis/dis2MbInNames.txt") as f:
        for line in f:
            name, cmd = line.split(maxsplit=1)
            input_cmd[name.replace("./", "")] = cmd.replace("./", "./dis/").strip()

    # Load task specifications
    with open(pathName+"/all_tasks.csv", "r") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter = ',')
        line = 0
        for row in csv_reader:
            if line > 0:
                all_tasks.append(row)
            line += 1 
    with open(pathName+"/levelAB_pairs.csv", "r") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter = ',')
        line = 0
        for row in csv_reader:
            if line > 0:
                levelAB.append(row)
            line += 1
    with open(pathName+"/levelC_threads.csv", "r") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter = ',')
        line = 0
        ccx = "0"
        threaded = True
        for row in csv_reader:
            if(len(row) < 3):
                continue
            # Start of a cluster - read definition
            if(row[0] == "threaded" or row[0] == "solo"):
                threaded = row[0] == "threaded"
                if(row[1] == "4"):
                    ccx = "1"
                else:
                    ccx = "0"
                continue
            # Start of a thread speficication
            if(row[0] != "task id" and row[0] != ""):
                row.append(ccx)
                if(threaded):
                    row.append("t")
                else:
                    row.append("s")
                levelC.append(row)
            line += 1
    ### Allocate Cache Ways ###
    run('mount -t resctrl resctrl /sys/fs/resctrl')
#    run('echo "L3:0=0000,1=0000,2=0000,3=0000" > /sys/fs/resctrl/schemata')
    with open(pathName+"/l3alloc.csv", "r") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter = ',')
        line = 0
        ccx = 0
        for row in csv_reader:
            if line > 0:
                core = int(row[0])
                way_AB = int(row[1])
                way_C = int(row[3])
                # Level C allocation
                if(core == 1):
                    run("mkdir -p /sys/fs/resctrl/level-c-ccx-0")
                    run('echo "L3:0=' + ("%0.4x" % (int("1"*way_C, 2))) + ';1=0000;2=0000;3=0000" > /sys/fs/resctrl/level-c-ccx-0/schemata')

                if(core == 4):
                    run("mkdir -p /sys/fs/resctrl/level-c-ccx-1")
                    run('echo "L3:0=0000;1=' + ("%0.4x" % (int("1"*way_C, 2))) + ';2=0000;3=0000" > /sys/fs/resctrl/level-c-ccx-1/schemata')
                run('mkdir -p /sys/fs/resctrl/level-ab-core-' + str(core))
                if(core < 4):
                    run('echo "L3:0=' + ("%0.4x" % (int("1"*way_AB + "0"*(16-way_AB), 2))) + ';1=0000;2=0000;3=0000" > /sys/fs/resctrl/level-ab-core-' + str(core) + '/schemata')
                else:
                    run('echo "L3:0=0000;1=' + ("%0.4x" % (int("1"*way_AB + "0"*(16-way_AB), 2))) + ';2=0000;3=0000" > /sys/fs/resctrl/level-ab-core-' + str(core) + '/schemata')
            line += 1
    ### Cleanup Old State ###
    run("rm -rf /dev/shm/*")
    pairID = 0
    ### Dispatch all the Level-A and -B paired and unpaired tasks ###
    for pairing in levelAB:
        task1 = int(pairing[P_T1_ID])
        task2 = int(pairing[P_T2_ID])
        # Lookup periods and binary location
        name1 = all_tasks[task1][A_BIN]
        name2 = all_tasks[task2][A_BIN]
        period1 = all_tasks[task1][A_PERIOD]
        period2 = all_tasks[task2][A_PERIOD]
        suite1 = all_tasks[task1][A_SUITE]
        suite2 = all_tasks[task2][A_SUITE]
        core = pairing[P_CORE]
        criticality = pairing[P_CRIT]
        binary = ""
        binary1 = ""
        binary2 = ""
        arg = ""
        # If the IDs match, this isn't a pair at all
        if task1 == task2:
            if suite1 == "TACLe":
                binary = "./baseline/bin/" + name1
            elif suite1 == "SD-VBS":
                binary = "./SD-VBS/" + name1 + "_qcif"
            elif suite1 == "DIS":
                binary = "./dis/" + name1
            else:
                print("Invalid suite name " + suite1 + " in Level-A/-B tasks! Exiting...")
                return 1

            lookupID = name1 + str(task1)
            if ID2PID(lookupID):
                print("ERROR: Task ID {} is already running! Did you end the previous case study?".format(lookupID))
                return 1
            # Arg format: <unique name> <infinite loops> <core> <NULL runID> <no output> <period> <crit lvl>
            arg = " " + lookupID + " -1 " + core + " NULL 0 " + period1 + " " + criticality
            bench_tsk = subprocess.Popen(input_cmd[name1] + " | numactl --interleave=all " + binary + arg, shell=True, executable='/bin/bash')
            pid_str = ID2PID(lookupID)
            if not pid_str:
                print("Unable to launch {} as a solo task! Exiting...".format(name1))
                return 1
            # Add this task to the appropriate cache partition (Class of Service in Intel CAT terms)
            run("echo " + pid_str + " > /sys/fs/resctrl/level-ab-core-" + core + "/tasks")
            addpid(pid_str)
        else:
            if suite1 == "TACLe":
                binary1 = "./baseline/bin/" + name1
                binary2 = "./baseline/bin/" + name2
            elif suite1 == "SD-VBS":
                binary1 = "./SD-VBS/" + name1 + "_qcif_all"
                binary2 = "./SD-VBS/" + name2 + "_qcif_all"
            elif suite1 == "DIS":
                binary1 = "./dis/" + name1 + "_all"
                binary2 = "./dis/" + name2 + "_all"
            else:
                print("Invalid suite name " + suite1 + " in Level-A/-B tasks! Exiting...")
                return 1

            pairID += 1 # Used for sem/shm identification in the pair task
            thread = str(int(core) + 16) # How Linux numbers SMT threads on the 3950X
            lookupID1 = name1 + str(task1)
            lookupID2 = name2 + str(task2)
            if ID2PID(lookupID1) or ID2PID(lookupID2):
                print("ERROR: Part of the pair containing task IDs {} and {} is already running! Did you end the previous case study?".format(lookupID1, lookupID2))
                return 1
            # Arg format: <unique name> <infinite loops> <core> <0 other core> <NULL other name> <NULL runID> <no output> <pairID> <period> <crit lvl>
            # XXX XXX XXX: FIXME when LITMUS is fixed
            arg1 = " " + lookupID1 + " -1 " + core + " NULL 0 " + period1 + " " + criticality
            arg2 = " " + lookupID2 + " -1 " + thread + " NULL 0 " + period2 + " " + criticality
            bench_tsk1 = subprocess.Popen(input_cmd[name1] + " | numactl --membind=0 " + binary1 + arg1, shell=True, executable='/bin/bash')
            bench_tsk2 = subprocess.Popen(input_cmd[name2] + " | numactl --membind=1 " + binary2 + arg2, shell=True, executable='/bin/bash')
            # Find benchmark PID (avoid getting shell or numactl PID)
            pid1_str = ID2PID(lookupID1)
            pid2_str = ID2PID(lookupID2)
            if not pid1_str or not pid2_str:
                print("Unable to launch {} and {} in a pair! Exiting...".format(name1, name2))
                return 1
            # Add this task to the appropriate cache partition (Class of Service in Intel CAT terms)
            run("echo " + pid1_str + " > /sys/fs/resctrl/level-ab-core-" + core + "/tasks")
            run("echo " + pid2_str + " > /sys/fs/resctrl/level-ab-core-" + core + "/tasks")
            addpid(pid1_str)
            addpid(pid2_str)

    ### Dispatch all the Level-C tasks ###
    for task in levelC:
        tid = int(task[T_ID])
        name = all_tasks[tid][A_BIN]
        period = all_tasks[tid][A_PERIOD]
        suite = all_tasks[tid][A_SUITE]
        ccx = task[T_CCX]

        if suite == "TACLe":
            binary = "./baseline/bin/" + name
        elif suite == "SD-VBS":
            name += "_qcif"
            binary = "./SD-VBS/" + name
        elif suite == "DIS":
            binary = "./dis/" + name
        else:
            print("Invalid suite name " + suite + " in Level-C tasks! Exiting...")
            return 1
        # XXX XXX XXX: FIXME when LITMUS is fixed
        core = str(int(ccx)*4 + 1);

        lookupID = name + str(tid)
        if ID2PID(lookupID):
            print("ERROR: Task ID {} is already running! Did you end the previous case study?".format(lookupID))
            return 1
        # Arg format: <unique name> <infinite loops> <core> <NULL runID> <no output> <period> <crit lvl>
        arg = " " + lookupID + " -1 " + core + " NULL 0 " + period + " " + str(CRIT_LEVEL_C)
        bench_tsk = subprocess.Popen(input_cmd[name] + " | numactl --interleave=all " + binary + arg, shell=True, executable='/bin/bash')
        # Find benchmark PID (avoid getting shell or numactl PID)
        pid_str = ID2PID(lookupID)
        if not pid_str:
            print("Unable to launch {} as a Level-C thread! Exiting...".format(name))
            return 1
        # Add this task to the appropriate cache partition (Class of Service in Intel CAT terms)
        run("echo " + pid_str + " > /sys/fs/resctrl/level-c-ccx-" + ccx + "/tasks")
        addpid(pid_str)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: " + sys.argv[0] + " <task sets path>")
        exit(1)
    if os.geteuid() != 0:
        print("You must run this script as root to enable cache isolation and access LITMUS-RT.")
        exit(1)
    pathName = sys.argv[1]
    err = main(pathName)
    input("Done. Waiting for termination signal...")
    exit(err)