import csv
import os
import sys
import tempfile
from pathlib import Path

import gradio as gr
import laspy
import numpy as np
import plotly.graph_objects as go

csv.field_size_limit(sys.maxsize)

POINT_CLOUD_CACHE = {}

# Define classification categories and colors
CLASSIFICATION_COLORS = {
    1: ("Soil", "saddlebrown"),
    2: ("Terrain", "peru"),
    3: ("Vegetation", "green"),
    4: ("Building", "red"),
    5: ("Street elements", "yellow"),
    6: ("Water", "blue"),
}


def load_point_cloud(las_path):
    if las_path in POINT_CLOUD_CACHE:
        return POINT_CLOUD_CACHE[las_path]

    with laspy.open(las_path) as fh:
        las = fh.read()
        point_count = len(las.points)

        x, y, z = las.x, las.y, las.z
        classification = las.classification
        intensity = las.intensity
        intensity = np.clip(np.log(intensity + 1) / 10.0, 0, 1)

        try:
            r, g, b = (
                las.red / 65535,
                las.green / 65535,
                las.blue / 65535,
            )
            has_color = True
        except:
            has_color = False
            r = g = b = None

    point_cloud_data = {
        "x": x,
        "y": y,
        "z": z,
        "intensity": intensity,
        "classification": classification,
        "has_color": has_color,
        "r": r,
        "g": g,
        "b": b,
        "point_count": len(x),
        "original_count": point_count,
    }

    POINT_CLOUD_CACHE[las_path] = point_cloud_data
    return point_cloud_data


def visualize_point_cloud(point_cloud_data, point_size=2, color_mode="RGB"):
    x, y, z = point_cloud_data["x"], point_cloud_data["y"], point_cloud_data["z"]
    legend_title = None
    fig_title_suffix = color_mode
    traces = []
    if color_mode == "RGB" and point_cloud_data["has_color"]:
        fig_title_suffix = "RGB"
        rgb_colors = [
            f"rgb({int(r*255)},{int(g*255)},{int(b*255)})"
            for r, g, b in zip(
                point_cloud_data["r"], point_cloud_data["g"], point_cloud_data["b"]
            )
        ]
        traces.append(
            go.Scatter3d(
                x=x,
                y=y,
                z=z,
                mode="markers",
                marker=dict(size=point_size, color=rgb_colors, opacity=0.8),
                name="Points (RGB)",
                showlegend=False,
            )
        )

    elif color_mode == "Intensity":
        fig_title_suffix = "Intensity"
        traces.append(
            go.Scatter3d(
                x=x,
                y=y,
                z=z,
                mode="markers",
                marker=dict(
                    size=point_size,
                    color=point_cloud_data["intensity"],
                    colorscale="Hot",  # Popular heat palettes: Hot, YlOrRd, Inferno, Magma
                    opacity=0.8,
                    colorbar=dict(
                        title="Intensity",
                        thickness=15,
                        len=0.75,
                        yanchor="middle",
                        y=0.5,
                    ),
                ),
                name="Points (Intensity)",
                showlegend=False,
            )
        )

    else:
        # Assign colors based on classification

        fig_title_suffix = "Classification"
        legend_title = "Classes"
        for class_value, info in CLASSIFICATION_COLORS.items():

            label, color = info
            mask = point_cloud_data["classification"] == class_value
            print(mask.shape)
            print(z.shape)
            x_c, y_c, z_c = x[mask], y[mask], z[mask]
            traces.append(
                go.Scatter3d(
                    x=x_c,
                    y=y_c,
                    z=z_c,
                    mode="markers",
                    marker=dict(size=point_size, color=color, opacity=0.8),
                    name=f"{label} ({class_value})",  # This name appears in the legend
                )
            )

    fig = go.Figure(data=traces)

    # Add a legend for classification colors
    fig.update_layout(
        scene=dict(
            xaxis=dict(visible=False, showgrid=False, zeroline=False, title=""),
            yaxis=dict(visible=False, showgrid=False, zeroline=False, title=""),
            zaxis=dict(visible=False, showgrid=False, zeroline=False, title=""),
            aspectmode="data",  # Crucial for correct 3D aspect ratio
        ),
        margin=dict(l=0, r=0, b=0, t=50),  # Adjust top margin for title
        title=dict(text=f"Point Cloud ({fig_title_suffix})", x=0.5, xanchor="center"),
        legend_title_text=legend_title,
        legend=dict(
            traceorder="normal",  # Or "reversed"
            itemsizing="constant",  # Makes legend markers a consistent size
            # You can also position the legend, e.g.:
            # x=1.05, y=1,
            # xanchor='left', yanchor='top'
        ),
    )

    return fig


def process_upload(file, point_size, color_mode):
    if file is None:
        return None

    file_path = Path(file.name)

    pc_data = load_point_cloud(file_path)
    # os.unlink(file_path)

    return visualize_point_cloud(pc_data, point_size, color_mode)


EXAMPLES_DIR = Path("examples")
EXAMPLES_DIR.mkdir(exist_ok=True)

example_files = [
    EXAMPLES_DIR / "Hillside.las",
    EXAMPLES_DIR / "Park.las",
    EXAMPLES_DIR / "Apartments.las",
]

for example in example_files:
    if not example.exists():
        for test_file in example_files:
            if test_file.exists():
                POINT_CLOUD_CACHE[str(example)] = POINT_CLOUD_CACHE.get(str(test_file))
                break

with gr.Blocks(title="Turin3D Dataset Visualizer") as app:
    gr.Markdown(
        """
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Turin3D Dataset</title>
    <style>
        body {
            font-family: 'Segoe UI', sans-serif;
            line-height: 1.6;
            margin: 0;
            padding: 0;
            color: #333;
        }
        header {
            background-color: #FFEE8C;
            color: white !important;
            padding: 20px;
            text-align: center;
        }
        section {
            padding: 20px;
            margin: 20px auto;
            color-scheme: light;
        }
        ul {
            padding-left: 20px;
        }
        body, p, ul, li, strong, code {
            color: #333;
        }
    </style>
</head>
<body>
    <header>
        <h1 style="color:white !important;">🏙️ Turin3D Dataset</h1>
        <p style="color:white !important;">Showcase of some portions of Turin3D dataset</p>
    </header>
</body>
</html>
"""
    )  # Updated main page title

    gr.Markdown(
        """
        ## About the Turin3D Dataset and Our Research
        This application showcases portions of the **Turin3D dataset**, which is introduced in our paper:
        **"Turin3D: Evaluating Adaptation Strategies under Label Scarcity in Urban LiDAR Segmentation with Semi-Supervised Techniques"**

        Our work presents Turin3D, a new large-scale aerial LiDAR dataset for 3D point cloud semantic segmentation. It covers approximately 1.43 km² in the city centre of Turin, Italy, and contains almost 70 million points.
        The dataset is specifically designed to address the challenges of label scarcity in urban environments. While the validation and test sets are manually annotated to ensure reliable evaluation of segmentation techniques, the training set is intentionally left unlabelled. This encourages the development and benchmarking of semi-supervised and self-supervised learning approaches, which are crucial when extensive ground truth annotations are not available.

        The full Turin3D dataset will be made publicly available to support further research in outdoor point cloud segmentation.
        For more details, please refer to our paper.
        """
    )

    gr.Markdown(
        """
        ## About this Interactive Visualizer
        This visualizer allows you to explore selected example point clouds from the Turin3D dataset.
        You can interact with the 3D data, switch between different color modes (RGB, Intensity, and Classification for ground truth labels on provided examples), and adjust the point size for visualization.

        **Important Note:** To ensure a smooth and responsive experience in this web application, the point clouds displayed here are **sub-sampled to a lower resolution** compared to the original high-density dataset. The example files are illustrative portions chosen for this demonstration.
        """
    )

    gr.Markdown(
        "--- \nUse the controls below to load an example or adjust the visualization."
    )

    with gr.Row():
        with gr.Column(scale=1):
            example_dropdown = gr.Dropdown(
                choices=[f.name for f in example_files], label="Example Point Clouds"
            )
            point_size = gr.Slider(
                minimum=1, maximum=10, value=2, step=1, label="Point Size"
            )

            color_mode = gr.Radio(
                choices=["RGB", "Intensity", "Classification"],
                value="RGB",
                label="Color Mode",
            )

            with gr.Row():
                example_btn = gr.Button("Load Example")

        with gr.Column(scale=2):
            plot_output = gr.Plot(label="Point Cloud Visualization")

    example_btn.click(
        fn=lambda example, size, mode: visualize_point_cloud(
            load_point_cloud(str(EXAMPLES_DIR / example)),
            point_size=size,
            color_mode=mode,
        ),
        inputs=[example_dropdown, point_size, color_mode],
        outputs=plot_output,
    )


if __name__ == "__main__":
    app.launch(share=False)