from gradio_client import Client, handle_file
from pathlib import Path
import gradio as gr
import numpy as np
from sklearn.cluster import KMeans
import trimesh
import plotly.graph_objects as go
import plotly.express as px
import polars as pl
from scipy.spatial import cKDTree

from constants import BLOCK_SIZES, LEGO_COLORS_RGB


def get_client() -> Client:
    return Client("TencentARC/InstantMesh")  # TODO: enable global client.


def generate_mesh(img: Path | str, seed: int = 42) -> str:
    """Takes a img (path or bytes) and returns a str (path) to the generated .obj-file"""

    client = get_client()
    result = client.predict(
        input_image=handle_file(img), do_remove_background=True, api_name="/preprocess"
    )
    result = client.predict(
        input_image=handle_file(result),
        sample_steps=75,
        sample_seed=seed,
        api_name="/generate_mvs",
    )
    result = client.predict(api_name="/make3d")
    obj_file = result[0]

    return obj_file


# ---- STEP 4: SELECT VOXEL SIZE ----
def voxelize(mesh_path: str | Path, resolution: str):
    resolution = {"Small (16)": 16, "Medium (32)": 32, "Large (64)": 64}[resolution]

    mesh = trimesh.load(mesh_path)
    bounds = mesh.bounds
    voxel_size = (bounds[1] - bounds[0]).max() / resolution  # pitch
    voxels = mesh.voxelized(pitch=voxel_size)
    colors = tree_knearest_colors(1, mesh, voxels)  # one is faster and good enough.
    mesh_state = {"voxels": voxels, "mesh": mesh, "colors": colors}

    return mesh_state


def build_scene(mesh, voxels):
    """Writes trimesh scene to .obj file"""
    voxels_mesh = voxels.as_boxes().apply_translation((1.5, 0, 0))
    scene = trimesh.Scene([mesh, voxels_mesh])
    scene.export("scene.obj")

    return "scene.obj"


# ---- STEP 5: VISUALIZE VOXELS ----
def quantize_colors(colors, k: int = 16):
    """
    quantize colors by fitting into 16 unique colors.
    """
    original_colors = np.array(colors)[:, :3]

    kmeans = KMeans(n_clusters=k, random_state=42)
    kmeans.fit(original_colors)

    # Get the representative colors
    representative_colors = kmeans.cluster_centers_.astype(int)

    # Transform the original colors to representative colors
    transformed_colors = representative_colors[kmeans.labels_]

    return transformed_colors


def lego_colors(colors):
    """
    quantize colors by fitting into 16 unique colors.
    """
    original_colors = np.array(colors)[:, :3]
    # Use scipy cdist to calculate euclidean distance between original and LEGO_C..
    from scipy.spatial.distance import cdist

    distances = cdist(original_colors, LEGO_COLORS_RGB, metric="sqeuclidean")
    distances = np.sqrt(distances)
    closest = np.argmin(distances, axis=1)

    return LEGO_COLORS_RGB[closest]


def pl_color_to_str():
    color_arr = pl.col("color").arr
    return pl.format(
        "rgb({},{},{})", color_arr.get(0), color_arr.get(1), color_arr.get(2)
    )


def visualize_voxels(mesh_state):
    # Step 1: Extract Colors
    # colors = tree_knearest_colors(5, mesh_state["mesh"], mesh_state["voxels"])
    # Step 2: Lego'ify Colors
    colors = mesh_state["colors"]
    # colors = quantize_colors(colors)
    # Step 3: Visualize
    voxels = mesh_state["voxels"]
    # Convert occupied_voxel_indices to a Polars DataFrame (if not already done)
    df = pl.from_numpy(voxels.sparse_indices, schema=["x", "z", "y"])
    df = df.with_columns(color=pl.Series(colors)).with_columns(
        color_str=pl_color_to_str()
    )

    return (
        px.scatter_3d(
            df,
            x="x",
            y="y",
            z="z",
            color="color_str",
            color_discrete_map="identity",
            symbol=["square"] * len(df),
            symbol_map="identity",
        ),
        df,
    )


def tree_knearest_colors(k: int, mesh, voxels):
    tree = cKDTree(mesh.vertices)
    distances, vertex_indices = tree.query(voxels.points, k=k)

    if k == 1:
        return mesh.visual.vertex_colors[vertex_indices]

    voxel_colors = []

    for nearest_indices in vertex_indices:
        neighbor_colors = mesh.visual.vertex_colors[nearest_indices]
        average_color = np.mean(neighbor_colors, axis=0).astype(np.uint8)
        voxel_colors.append(average_color)

    return voxel_colors


# ---- STEP 6: ADJUST BRIGHTNESS ----
# def adjust_brightness(image, brightness):
# adjusted_image = cv2.convertScaleAbs(image, alpha=brightness)
# return adjusted_image


# ---- STEP 8: LEGO BUILD ANIMATION ----
def merge_into_bricks(grouped_df: pl.DataFrame, BLOCK_SIZES) -> pl.DataFrame:
    color_str = grouped_df[0, "color_str"]
    z_val = grouped_df[0, "z"]

    xy_grid = np.zeros(
        (grouped_df["x"].max() + 1, grouped_df["y"].max() + 1), dtype=bool
    )
    xy_grid[grouped_df["x"], grouped_df["y"]] = 1
    out_rows = []
    grouped_df = grouped_df.sort(by=["x", "y"])
    coords = {(x, y) for x, y in grouped_df[["x", "y"]].to_numpy()}

    while coords:
        (x0, y0) = coords.pop()
        coords.add((x0, y0))  # reinsert until placed

        placed = False
        for width, height in BLOCK_SIZES:
            if x0 + width > xy_grid.shape[0] or y0 + height > xy_grid.shape[1]:
                continue
            if np.all(xy_grid[x0 : x0 + width, y0 : y0 + height] == 1):
                place_block(x0, y0, width, height, coords)
                xy_grid[x0 : x0 + width, y0 : y0 + height] = 0  # remove from xygrid
                out_rows.append((color_str, z_val, x0, y0, width, height))
                placed = True
                break

        if not placed:
            # fallback to 1x1
            coords.remove((x0, y0))
            out_rows.append((color_str, z_val, x0, y0, 1, 1))

    return pl.DataFrame(
        {
            "color_str": [row[0] for row in out_rows],
            "z": [row[1] for row in out_rows],
            "x": [row[2] for row in out_rows],
            "y": [row[3] for row in out_rows],
            "width": [row[4] for row in out_rows],
            "height": [row[5] for row in out_rows],
        }
    )


def can_place_block(x0, y0, w, h, coords):
    for xx in range(x0, x0 + w):
        for yy in range(y0, y0 + h):
            if (xx, yy) not in coords:
                return False
    return True


def place_block(x0, y0, w, h, coords):
    for xx in range(x0, x0 + w):
        for yy in range(y0, y0 + h):
            coords.remove((xx, yy))


# Function to generate vertices for a rectangular prism (brick)
def create_brick(x, y, z, width, height, depth=1, color="gray"):
    return go.Mesh3d(
        x=[x, x + width, x + width, x, x, x + width, x + width, x],  # X-coordinates
        y=[y, y, y + height, y + height, y, y, y + height, y + height],  # Y-coordinates
        z=[z, z, z, z, z + depth, z + depth, z + depth, z + depth],  # Z-coordinates
        color=color,
        alphahull=-1,
        i=[7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2],
        j=[3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3],
        k=[0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6],
        name=f"Z={z}",
    )


def get_range(series: pl.Series) -> tuple[int, int]:
    return series.min(), series.max()


def animate_lego_build(df_state):
    # Colors already merged.
    df: pl.DataFrame = df_state
    df = df.with_columns(color=quantize_colors(df["color"])).with_columns(
        color_str=pl_color_to_str()
    )
    # Quantize Colors... Need to split string and use..
    merged_df = df.group_by("color_str", "z").map_groups(
        lambda grp: merge_into_bricks(grp, BLOCK_SIZES)
    )

    fig = go.Figure()
    fig.update_layout(
        scene=dict(
            xaxis=dict(range=get_range(df["x"]), autorange=False),
            yaxis=dict(range=get_range(df["y"]), autorange=False),
            zaxis=dict(range=get_range(df["z"]), autorange=False),
        )
    )

    # Add each brick to the plot
    for z in merged_df["z"].unique().sort():
        for row in merged_df.filter(pl.col("z") == z).iter_rows(named=True):
            fig.add_trace(
                create_brick(
                    x=row["x"],
                    y=row["y"],
                    z=row["z"],
                    width=row["width"],
                    height=row["height"],
                    color=row["color_str"],
                )
            )
        # frame_jpgs.append(f"frame_z_{z}.jpg")
        # if not Path(frame_jpgs[-1]).exists():
        #    fig.write_image(frame_jpgs[-1])

    return fig  # , frame_jpgs


# ---- GRADIO UI ----
with gr.Blocks() as demo:
    gr.Markdown("# 🧱 **Image 2 Lego Builder** 🧱")

    # Step 1: Upload Image and Build Mesh
    with gr.Column(variant="compact"):
        with gr.Row():
            image_input = gr.Image(
                type="filepath", height="250px", label="Upload an Image"
            )
            with gr.Column(variant="compact"):
                seed = gr.Number(label="Seed", value=42)
                # Potentially add color options.
                voxel_size_selector = gr.Dropdown(
                    ["Small (16)", "Medium (32)", "Large (64)"],
                    value="Medium (32)",
                    label="Select Voxel Size",
                )
        with gr.Row():
            build_button = gr.Button("Generate Mesh")
            voxelize_button = gr.Button("Generate Voxels")

    # Visualizations...
    # Mesh | Voxel Color | Voxel Lego Bricks+Color
    with gr.Row():
        mesh_info_display = gr.Model3D(
            label="Mesh Visualization", height="250px", value="mesh.obj"
        )
        voxel_color_display = gr.Plot(label="Colorized Voxels")
        voxel_bricks = gr.Plot(label="Lego Bricks")
        brick_animation = gr.Gallery(label="Build Animation")

    mesh_state = gr.State(value={})

    build_button.click(
        generate_mesh, inputs=[image_input, seed], outputs=mesh_info_display
    )

    # Step 4: Select Voxel Size
    voxelize_button.click(
        voxelize,
        inputs=[mesh_info_display, voxel_size_selector],
        outputs=[mesh_state],
    )
    df_state = gr.State()
    mesh_state.change(
        visualize_voxels,
        inputs=[mesh_state],
        outputs=[voxel_color_display, df_state],
    )

    df_state.change(animate_lego_build, inputs=[df_state], outputs=[voxel_bricks])

    def anim_pltly(df):
        df = df.with_columns(color=quantize_colors(df["color"])).with_columns(
            color_str=pl_color_to_str()
        )
        # Quantize Colors... Need to split string and use..
        merged_df = df.group_by("color_str", "z").map_groups(
            lambda grp: merge_into_bricks(grp, BLOCK_SIZES)
        )

        fig = go.Figure()
        fig.update_layout(
            scene=dict(
                xaxis=dict(range=get_range(df["x"]), autorange=False),
                yaxis=dict(range=get_range(df["y"]), autorange=False),
                zaxis=dict(range=get_range(df["z"]), autorange=False),
            )
        )
        frame_jpgs = []
        # Add each brick to the plot
        for z in merged_df["z"].unique().sort():
            for row in merged_df.filter(pl.col("z") == z).iter_rows(named=True):
                fig.add_trace(
                    create_brick(
                        x=row["x"],
                        y=row["y"],
                        z=row["z"],
                        width=row["width"],
                        height=row["height"],
                        color=row["color_str"],
                    )
                )
            frame_jpgs.append(f"frame_z_{z}.jpg")
            if not Path(frame_jpgs[-1]).exists():
                fig.write_image(frame_jpgs[-1])

        return frame_jpgs

    # TODO: add to generate layer-by-layer
    # df_state.change(anim_pltly, inputs=[df_state], outputs=[brick_animation])

# Launch the app
demo.launch(share=True, debug=True)