import gradio as gr
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
import tempfile
import traceback
import os

from types import SimpleNamespace

DEBUG = False # True: skip gradio and execute in main function

def process_point_cloud(file_obj):
    # Save uploaded file temporarily
    if DEBUG:
        temp_path = os.path.join(tempfile.gettempdir(), file_obj.name)
        with open(temp_path, "wb") as f:
            f.write(file_obj.read())
    else:
        temp_path = file_obj.name

    # Load point cloud using Open3D
    pcd = o3d.io.read_point_cloud(temp_path)

    # Process: voxel downsampling
    downsampled = pcd.voxel_down_sample(voxel_size=0.01)

    # Estimate normals for visualization
    downsampled.estimate_normals()

    points = np.asarray(downsampled.points)
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.5, c='b')
    ax.axis('off')
    plt.tight_layout()
    temp_img_path = os.path.join(tempfile.gettempdir(), "rendered.png")
    plt.savefig(temp_img_path, dpi=150)
    plt.close()

    img_np = plt.imread(temp_img_path)
    return img_np

iface = gr.Interface(
    fn=process_point_cloud,
    inputs=gr.File(file_types=[".ply", ".pcd", ".xyz"], label="Upload Point Cloud"),
    outputs=gr.Image(type="numpy", label="Rendered Point Cloud"),
    title="Point Cloud Viewer",
    description="Upload a .ply, .pcd, or .xyz file. The app will downsample and render it."
)

def check_delimiter(file_path):
    with open(file_path, "r") as f:
        line = f.readline()
    if "," in line:
        return ","
    elif "\t" in line:
        return "\t"
    else:
        return " "

def meshing(file_obj, normal_rad_coef, normal_max_nn, orient_k, bpa_min_coef, bpa_max_coef):
    try:
        if DEBUG:
            temp_path = os.path.join(file_obj.dir, file_obj.name)
            ext = os.path.splitext(temp_path)[1].lower()
        else:
            temp_path = file_obj.name
            ext = os.path.splitext(temp_path)[1].lower()

        # Load the input point cloud
        if ext == ".txt":
            delimiter = check_delimiter(temp_path)
            points = np.loadtxt(temp_path, delimiter=delimiter)
            if points.shape[1] != 3:
                raise ValueError("Text file must contain 3 columns (X Y Z)")
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(points)
        else:
            pcd = o3d.io.read_point_cloud(temp_path)

        dists = pcd.compute_nearest_neighbor_distance()
        avg_spacing = np.mean(dists)
        normal_radius = avg_spacing*normal_rad_coef # start with 2–3x spacing
        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=normal_radius, max_nn=normal_max_nn))
        pcd.orient_normals_consistent_tangent_plane(k=orient_k)

        # Mesh using Ball Pivoting
        coeffs = np.linspace(bpa_min_coef, bpa_max_coef, int(4))
        mesh_radii = [avg_spacing * x for x in coeffs]
        mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(pcd, o3d.utility.DoubleVector(mesh_radii))

        # Mesh using Poisson
        # mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9)

        # Export as OBJ
        mesh_output_path = os.path.join(tempfile.gettempdir(), "mesh_output.obj")
        o3d.io.write_triangle_mesh(mesh_output_path, mesh, write_ascii=True)

        # Render image preview
        verts = np.asarray(mesh.vertices)
        tris = np.asarray(mesh.triangles)
        fig = plt.figure(figsize=(6, 6))
        ax = fig.add_subplot(111, projection='3d')
        ax.plot_trisurf(verts[:, 0], verts[:, 1], verts[:, 2], triangles=tris, color='lightblue', edgecolor='gray', linewidth=0.1)
        ax.axis('off')
        plt.tight_layout()
        image_path = os.path.join(tempfile.gettempdir(), "mesh_render.png")
        plt.savefig(image_path, dpi=150)
        plt.close()
        img_np = plt.imread(image_path)

        return img_np, mesh_output_path

    except Exception as e:
        traceback.print_exc()
        return np.zeros((200, 400, 3), dtype=np.uint8), None

iface_mesh = gr.Interface(
    fn=meshing,
    inputs=[
        gr.File(file_types=[".txt", ".ply", ".pcd"], label="Upload Point Cloud"),
        gr.Slider(1.0, 10.0, value=2.0, step=0.1, label="Normal Radius (Hybrid Search)"),
        gr.Slider(1, 100, value=30, step=1, label="Max NN for Normals"),
        gr.Slider(1, 100, value=30, step=1, label="K for Consistent Normal Orientation"),
        gr.Slider(0.1, 10.0, value=1.5, step=0.1, label="Ball Pivoting Radius Min"),
        gr.Slider(0.1, 10.0, value=3.5, step=0.1, label="Ball Pivoting Radius Max"),
    ],
    outputs=[
        gr.Image(type="numpy", label="Mesh Preview"),
        gr.File(label="Download Mesh (STL)")
    ],
    title="Point Cloud Meshing to STL",
    description="Upload a point cloud (.txt, .ply, .pcd). The app creates a triangle mesh and exports it as an STL file."
)

if __name__ == "__main__":
    if DEBUG:
        test_file_path = "./samples/Panel_pca_by20_space.txt"
        with open(test_file_path, "rb") as f:
            dummy = SimpleNamespace()
            dummy.name = os.path.basename(test_file_path)
            dummy.dir = os.path.dirname(test_file_path)
            dummy.read = lambda: f.read()
            # result_img = process_point_cloud(dummy)
            result_img, result_mesh = meshing(dummy)
        out_path = os.path.join("temporal", "debug_render.png")
        plt.imsave(out_path, result_img)
        print(f"[DEBUG] Rendered image saved at: {out_path}")
    else:
        # iface.launch()
        iface_mesh.launch()