#!/usr/bin/python3
import json
import sys

import pycuda.driver as cuda
from gevo import irind

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Stage file name is needed!")
        exit(-1)

    stageFileName = sys.argv[1]

    llvm_src_filename='cuda-device-only-kernel.ll'
    try:
        with open(llvm_src_filename, 'r') as f:
            initSrcEnc = f.read().encode()
    except IOError:
        print("File {} does not exist".format(llvm_src_filename))
        exit(1)

    stage = json.load(open(stageFileName))

    pop = {}
    for entry in stage:
        pop[tuple(entry['fitness'])] = irind.encode_edits_from_list(entry['edits'])

    allfit = [tuple(entry['fitness']) for entry in stage]
    allfit = list(set(allfit))

    sort_by_perf = sorted(allfit, key=lambda x: x[0], reverse=True)
    sort_by_error = sorted(allfit, key=lambda x: x[1], reverse=True)

    sort_type = int(input("Sorted by [1]runtime or [2]error? "))
    if sort_type == 1:
        for i in sort_by_perf:
            print(i)
    else:
        for i in sort_by_error:
            print(i)

    desc = input("Dump which variant? ")
    key = eval(desc)

    # Detect GPU property
    cuda.init()
    # TODO: check if there are multiple GPUs.
    SM_MAJOR, SM_MINOR = cuda.Device(0).compute_capability()
    mgpu = 'sm_' + str(SM_MAJOR) + str(SM_MINOR)

    ind = irind.llvmIRrep(initSrcEnc, mgpu, pop[key])
    with open("dump.ll", 'w') as f:
        f.write(ind.srcEnc.decode())
    with open("dump.edit", 'w') as f:
        print(ind.edits, file=f)
    ind.ptx('gevo.ptx')
    print("Generate dump.ll and gevo.ptx")
