import numpy as np
from math import *
import pickle
import sys
import sdfGenerate
# sys.path.append('/home/dtouch/projects/implicit/implicit_rendering/renderSDF/sdf_generate/cmake-build-debug')

print('run test_make_sphere_sdf.py...')

sample_cnt_x, sample_cnt_y, sample_cnt_z = 188, 188, 188

spatial_size = (4, 4, 4)

rod_crystal_sdf = sdfGenerate.generate(sample_cnt_x, sample_cnt_y, sample_cnt_z, spatial_size[0]/2)


sdf = np.ones((sample_cnt_x, sample_cnt_y, sample_cnt_z))


def sd_plane_x(x, y, z, plane_pos):
    return abs(x - plane_pos)


def sd_plane_y(x, y, z, plane_pos):
    return abs(y - plane_pos)


def sd_plane_z(x, y, z, plane_pos):
    return abs(z - plane_pos)


def sd_box(p):
    def length(vec):
        return sqrt(vec[0] ** 2 + vec[1] ** 2 + vec[2] ** 2)
    size = np.array([40, 6, 18])
    center = np.array([30, 30, 30])
    q = abs(p - center) - size / 2
    return length(np.maximum(q, 0)) + min(np.amax(q), 0)


def sd_sphere(x, y, z):
    sphere_center = (6, 6, 6)
    radius = 5
    return sqrt((x - sphere_center[0]) ** 2 + (y - sphere_center[1]) ** 2 + (z - sphere_center[2]) ** 2) - radius


for i in range(sample_cnt_x):
    for j in range(sample_cnt_y):
        for k in range(sample_cnt_z):
            x = (spatial_size[0] * i / (sample_cnt_x - 1))
            y = (spatial_size[1] * j / (sample_cnt_y - 1))
            z = (spatial_size[2] * k / (sample_cnt_z - 1))
            sdf[i][j][k] = sd_sphere(x, y, z)
            # sdf[i][j][k] = sd_box(np.array([x, y, z]))
            # sdf[i][j][k] = sd_plane_x(x, y, z, 0)

sdf = rod_crystal_sdf

# save_file = '/home/dtouch/projects/implicit/implicit_rendering/renderSDF/data/sphere.sdf'
# save_file = '/home/dtouch/projects/implicit/implicit_rendering/renderSDF/data/box.sdf'
save_file = '/home/dtouch/projects/implicit/implicit_rendering/renderSDF/data/crystal_rod.sdf'

with open(save_file, 'wb') as f:
    pickle.dump([sdf], f)

data = pickle.load(open(save_file, 'rb'))