import sys
from pathlib import Path

import numpy as np
from pyglm.glm import vec3
from pyxpg import AllocType, Buffer, BufferUsageFlags, Format, Image, ImageLayout, ImageUsageFlags, MemoryUsage, imgui

from ambra.config import Config, GuiConfig
from ambra.lights import DirectionalLight, DirectionalShadowSettings
from ambra.marching_cubes import MarchingCubesPipeline
from ambra.primitives3d import Grid, GridType, MarchingCubesMesh, Mesh
from ambra.utils.gpu import readback_buffer, view_bytes
from ambra.utils.hook import hook
from ambra.viewer import Viewer

if False:
    x, y, z = 128, 128, 128
    sdf = np.zeros((x, y, z), np.float32)
    size = (1.0, 1.0, 1.0)
else:
    vol_path = Path(sys.argv[1])
    sdf: np.ndarray = np.load(vol_path)["volume"]
    size = np.array(sdf.shape[::-1], np.float32) / max(sdf.shape)
    z, y, x = sdf.shape

    if True:
        # (A, Z, Y, X, C)
        sdf = sdf.reshape((1, z, y, x, 1))

        N = 32
        sdf = np.repeat(sdf, N, axis=0)
        for i in range(N):
            sdf[i] += i / N

    if False:
        # Dragon big
        # 1662892 positions
        # 7638876 primitives
        #
        # 1654432 area
        #
        # expected_positions = max(area * 2, 1024)
        # expected_triangles = max(area * 10, 5 * 1024)
        #
        # Dragon small
        #  65248
        # 300216
        #
        # 102634 area
        area = (x * y + x * z + y * z) * 2
        print(x, y, z, area)


class CustomViewer(Viewer):
    @hook
    def on_gui(self):
        global m
        if isinstance(m, MarchingCubesMesh):
            if imgui.begin("MC")[0]:
                u, m.level = imgui.slider_float("Level", m.level, -5, 5)
                if u:
                    m.invalidate_surface()
            imgui.end()


v = CustomViewer(config=Config(gui=GuiConfig(stats=True)))

print(v.device.features)
print(v.device.subgroup_size_control)
print(v.device.compute_full_subgroups)
print(v.device.device_properties.subgroup_size_control_properties.min_subgroup_size)
print(v.device.device_properties.subgroup_size_control_properties.max_subgroup_size)
print(v.device.device_properties.subgroup_size_control_properties.max_compute_workgroup_subgroups)
print(v.device.device_properties.subgroup_size_control_properties.required_subgroup_size_stages)

if False:
    # print(v.renderer.device.device_properties.subgroup_properties.subgroup_size)
    pipeline = MarchingCubesPipeline(v.renderer)

    sdf_buf = Buffer.from_data(
        v.device, view_bytes(sdf), BufferUsageFlags.TRANSFER_SRC, AllocType.HOST, name="sdf-buf"
    )
    sdf_gpu = Image(
        v.device,
        x,
        y,
        depth=z,
        format=Format.R32_SFLOAT,
        usage_flags=ImageUsageFlags.STORAGE | ImageUsageFlags.TRANSFER_DST,
        alloc_type=AllocType.DEVICE,
        name="sdf",
    )

    with v.device.sync_commands() as cmd:
        cmd.image_barrier(sdf_gpu, ImageLayout.TRANSFER_DST_OPTIMAL, MemoryUsage.NONE, MemoryUsage.TRANSFER_DST)
        cmd.copy_buffer_to_image(sdf_buf, sdf_gpu)
        cmd.image_barrier(sdf_gpu, ImageLayout.GENERAL, MemoryUsage.TRANSFER_DST, MemoryUsage.COMPUTE_SHADER)

    res = pipeline.run_sync(v.renderer, sdf_gpu, size, 0.0)

    positions = readback_buffer(v.device, res.positions).view(np.float32).reshape((-1, 3))
    normals = readback_buffer(v.device, res.normals).view(np.float32).reshape((-1, 3))
    indices = readback_buffer(v.device, res.indices).view(np.uint32)

    print(positions.shape)
    print(normals.shape)
    print(indices.shape)
    m = Mesh(positions, indices, normals)
else:
    m = MarchingCubesMesh(sdf[:, :, :], size)

l1 = DirectionalLight.look_at(
    vec3(5, 6, 7) * 4.0,
    vec3(0, 0, 0),
    vec3(0, 0, 1),
    np.array([1.0, 1.0, 1.0]),
    shadow_settings=DirectionalShadowSettings(half_extent=10.0, z_near=1.0, z_far=100),
)
l2 = DirectionalLight.look_at(
    vec3(-5, -6, 7) * 4.0,
    vec3(0, 0, 0),
    vec3(0, 0, 1),
    np.array([1.0, 1.0, 1.0]),
    shadow_settings=DirectionalShadowSettings(half_extent=10.0, z_near=1.0, z_far=100),
)

g = Grid.transparent_black_lines((100, 100), GridType.XY_PLANE)

v.scene.objects.extend([m, l1, l2, g])
v.run()
