import argparse
import os
import subprocess
from pathlib import Path

import numpy as np
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument("--executable", type=str, default="bin/surface_renderer")
parser.add_argument("--input-mesh-path", type=str, required=True)
args = parser.parse_args()

output_npz_path = "view_images.npz"
image_output_directory = Path("images")
num_viewpoints = 2
image_width = 400
image_height = 400
command = [
    args.executable, "--input-mesh-path", args.input_mesh_path,
    "--output-npz-path", output_npz_path, "--num-viewpoints", num_viewpoints,
    "--image-width", image_width, "--image-height", image_height
]
command = [str(arg) for arg in command]
result = subprocess.run(command,
                        encoding="utf-8",
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE)
print(result.stdout)
print(result.stderr)

data = np.load(output_npz_path)
print(data.files)

os.makedirs(image_output_directory, exist_ok=True)
for key in data.files:
    image = data[key]
    image = image.reshape((image_height, image_width, 3))
    image = (image + 1) / 2 * 255
    image = image.astype(np.uint8)
    image = Image.fromarray(image)
    image.save(image_output_directory / f"{key}.png")
