import numpy as np
import pickle
import math

def padding(filePath, outputFile):
    data = pickle.load(open(filePath, 'rb'))
    x, y, z = data[0].shape[0], data[0].shape[1], data[0].shape[2]
    print(x, y, z)
    size = np.power(2, math.ceil(math.log2(x)))
    data = data[0]
    size = max(size, np.power(2, math.ceil(math.log2(y))))
    size = max(size, np.power(2, math.ceil(math.log2(z))))
    print('size = ', size)
    pad_width = ((0, size-x), (0, size-y), (0, size-z))
    data = np.pad(data, pad_width, mode='constant', constant_values=-999)
    print(data.shape)
    saveArrayToRaw(data, outputFile)

def cut(filePath, outputFile):
    data = pickle.load(open(filePath, 'rb'))
    x, y, z = data[0].shape[0], data[0].shape[1], data[0].shape[2]
    size = np.power(2, math.floor(math.log2(x)))
    
    size = min(size, np.power(2, math.floor(math.log2(y))))
    size = min(size, np.power(2, math.floor(math.log2(z))))
    print('size = ', size)
    data = data[0][:size, :size, :size].flatten()
    saveArrayToRaw(data, outputFile)


def loadSdfFile(filePath):
    data = pickle.load(open(filePath, 'rb'))
    x, y, z = data[0].shape[0], data[0].shape[1], data[0].shape[2]
    print(x, y, z)
    data = data[0].flatten()
    print(data)
    return x, y, z, data

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

def loadRawFile(filePath):
    data = np.fromfile(filePath, dtype=np.float32)
    print(data)
    return data



def check():
    x, y, z, data = loadSdfFile('data/mesh_test_step_4.sdf')
    saveArrayToRaw(data, './data.raw')
    new_data = loadRawFile('./data.raw')
    cnt = 0
    for i in range(x):
        for j in range(y):
            for k in range(z):
                if cnt == 11675:
                    print(i, j, k)
                    exit(-1)
                cnt = cnt + 1
                idx = i * y * z + j * z + k
                if math.fabs(data[idx] - new_data[idx]) > 1e-5:
                    print("error: ", data[idx], '---', new_data[idx])
                    exit(-1)
    print("check success")


if __name__ == '__main__':
    #check()
    padding('data/bunny_test_step_0.000312.sdf', './data/bunny_test.raw')