import os
import pickle
import sys
import math
from optparse import OptionParser
import numpy as np
import test_make_sphere_sdf

def init():
    usage = "usage: python3 run.py -f inputFile\n(need to install numpy)"
    parser = OptionParser(usage)
    parser.add_option("-f", "--file", dest="file",
                      help="read data from FILENAME(need .raw file)")
    parser.add_option("-d", "--demo", dest="demo",
                      help="Implicit function demo")
    parser.add_option("-s", "--size", dest="size",
                      help="Rendering resolution, must be a power of 2, eg: 128, 256")
    # parser.add_option("-s", "--size", dest="dataSize",
    #                   help="the size of input data(x, y, z)")
    (optionsMap, argsList) = parser.parse_args()
    if len(sys.argv) <= 1:
        parser.print_help()
        exit(-1)
    return optionsMap, argsList


# data is a 3D array
def loadSdfFile(file_path):
    data = pickle.load(open(file_path, 'rb'))
    x, y, z = data[0].shape[0], data[0].shape[1], data[0].shape[2]
    data = data[0]
    print("data:")
    print(data)
    print("x,", x, "y,", y, "z,", z)
    data = -data
    return x, y, z, data


def saveArrayToRaw(array, output_file, element_type='float32'):
    new_array = array.astype(element_type)
    new_array.tofile(output_file)


def run(options, pack_value=-999):
    status = 0
    if os.path.exists("./build"):
        status = os.system("cd build && cmake ..")
    else:
        status = os.system("mkdir build && cd build && cmake ..")
    # compile
    if status != 0:
        os._exit(255)
    os.chdir("build")
    os.system("make")
    os.chdir("../")

    # padding and saving raw files
    file = options.file
    if file:
        file_name, file_ext = os.path.splitext(file)
        print(file_ext)
        size = 256
        if file_ext == '.sdf':
            x, y, z, data = loadSdfFile(file)
            size = np.power(2, math.ceil(math.log2(x)))
            size = max(size, np.power(2, math.ceil(math.log2(y))))
            size = max(size, np.power(2, math.ceil(math.log2(z))))
            pad_width = ((0, size-x), (0, size-y), (0, size-z))
            data = np.pad(data, pad_width, mode='constant', constant_values=pack_value)
            print(data.shape)
            data = data.flatten()
            file = file_name + ".raw"
            saveArrayToRaw(data, file)

            # exec
            print("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -file " + file + " -size " + str(math.log2(size)))
            os.system("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -file " + file + " -size " + str(math.log2(size)))
        elif file_ext == '.sdf':

            # exec
            print("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -file " + file + " -size " + str(math.log2(size)))
            os.system("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -file " + file)
        else:
            raise Exception('Input file extension error!')
    else:
        render_type = options.demo
        size = 128
        if options.size:
            size = int(options.size)
        print("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -type " + render_type + " -size " + str(math.log2(size)))
        os.system("__NV_PRIME_RENDER_OFFLOAD=1 __GLX_VENDOR_LIBRARY_NAME=nvidia ./build/renderSDF -type " + render_type + " -size " + str(math.log2(size)))


if __name__ == "__main__":
    opt, arg = init()
    run(opt)