import numpy as np


def depth_to_pointcloud(depth, image=None, stride=4):
    """
    Converte um mapa de profundidade em uma nuvem de pontos 3D.

    depth:
        Array 2D com os valores de profundidade.

    image:
        Imagem RGB original para colorir os pontos.

    stride:
        Intervalo entre pixels.
        stride=4 significa que usamos 1 a cada 4 pixels.
    """

    depth = np.asarray(depth, dtype=np.float32)

    if depth.ndim != 2:
        raise ValueError(
            f"Depth precisa ser 2D. Formato recebido: {depth.shape}"
        )

    height, width = depth.shape

    # Seleciona os pixels
    y, x = np.mgrid[
        0:height:stride,
        0:width:stride
    ]

    z = depth[::stride, ::stride]

    # Remove valores inválidos
    valid = np.isfinite(z)

    x = x[valid].astype(np.float32)
    y = y[valid].astype(np.float32)
    z = z[valid].astype(np.float32)

    if len(z) == 0:
        raise ValueError(
            "Não foram encontrados valores válidos de profundidade."
        )

    # Normaliza Z entre 0 e 1
    z_min = np.min(z)
    z_max = np.max(z)

    if z_max - z_min > 1e-8:
        z = (z - z_min) / (z_max - z_min)
    else:
        z[:] = 0.5

    # Centraliza X e Y
    x = (x - width / 2) / width
    y = -(y - height / 2) / height

    # Aumenta a profundidade para facilitar a visualização
    z = z * 2.0

    points = np.stack(
        [x, y, z],
        axis=1
    )

    colors = None

    if image is not None:

        image = image.convert("RGB")

        image_array = np.asarray(
            image.resize((width, height))
        )

        sampled_colors = image_array[
            ::stride,
            ::stride
        ]

        colors = sampled_colors[valid]

    return points, colors


def save_pointcloud_ply(filename, points, colors=None):
    """
    Salva uma nuvem de pontos no formato PLY.
    """

    points = np.asarray(points)

    if points.ndim != 2 or points.shape[1] != 3:
        raise ValueError(
            "points precisa ter formato (N, 3)."
        )

    with open(filename, "w", encoding="utf-8") as file:

        file.write("ply\n")
        file.write("format ascii 1.0\n")
        file.write(f"element vertex {len(points)}\n")

        file.write("property float x\n")
        file.write("property float y\n")
        file.write("property float z\n")

        if colors is not None:

            file.write("property uchar red\n")
            file.write("property uchar green\n")
            file.write("property uchar blue\n")

        file.write("end_header\n")

        if colors is not None:

            for point, color in zip(points, colors):

                x, y, z = point
                r, g, b = color

                file.write(
                    f"{x:.6f} "
                    f"{y:.6f} "
                    f"{z:.6f} "
                    f"{int(r)} "
                    f"{int(g)} "
                    f"{int(b)}\n"
                )

        else:

            for point in points:

                x, y, z = point

                file.write(
                    f"{x:.6f} "
                    f"{y:.6f} "
                    f"{z:.6f}\n"
                )