import numpy as np
from pypsdr.plot import write_obj
from pypsdr.utils.io import imread, imwrite
import cv2

# parameters
scale = 0.02
size = 10


def getMesh(map, size):
    res = map.shape[0]
    y = map
    x, z = np.meshgrid(np.linspace(-size, size, res),
                       np.linspace(-size, size, res))
    idx = np.arange(res * res).reshape((res, res))
    V = np.stack([x.flatten(), y.flatten(), z.flatten()], axis=1)
    F = np.stack([idx[:-1, :-1].flatten(), idx[:-1, 1:].flatten(),
                  idx[1:, 1:].flatten(), idx[1:, :-1].flatten()], axis=1)
    F = F[:, ::-1]
    return V, F


# read height map
image = imread("map.png")[:, :, 0]

height_map = cv2.resize(image, (128, 128)) * scale
height_map = cv2.resize(height_map, (256, 256))
V, F = getMesh(height_map, size)

write_obj(V, F, "tar.obj")

init_map = cv2.blur(image, (256, 256)) * scale
init_map = cv2.resize(init_map, (256, 256))
V_init, F_init = getMesh(init_map, size)

write_obj(V_init, F_init, "init.obj")
