#!/usr/bin/env python

"""
FingerNet Demo App
"""

import os
import numpy as np
import trimesh
import pandas as pd
import tempfile
import gradio as gr
import matplotlib as mpl
from huggingface_hub import hf_hub_download
import onnxruntime as ort

model_dict = {
    "FingerNet": "fingernet",
    "FingerNet-Surf": "fingernet-surf",
}

vertices_dict = {}
faces_dict = {}
def_index_dict = {}
for name in model_dict.keys():
    vertices_dict[name] = np.loadtxt(
        f"assets/{name.lower().replace('net', '')}/surf_coor.txt", delimiter=","
    )
    faces_dict[name] = (
        np.loadtxt(
            f"assets/{name.lower().replace('net', '')}/surf_tri.txt",
            delimiter=",",
            dtype=int,
        )
        - 1
    )
    def_index_dict[name] = np.loadtxt(
        f"assets/{name.lower().replace('net', '')}/def_node.txt", dtype=int
    )
colormap = mpl.colormaps.get_cmap("viridis")
cmin = 0.0
cmax = 12.0

def load_model(model_name="fingernet"):
    repo_id = f"asRobotics/{model_name}"
    model_filename = "model.onnx"
    model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
    return ort.InferenceSession(model_path)

ort_dict = {}
for name, path in model_dict.items():
    try:
        ort_dict[name] = load_model(name)
    except Exception as e:
        print(f"Warning: Could not load model '{name}' from '{path}'. Error: {e}")


def make_double_sided(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    faces_front = mesh.faces.copy()
    vertices = mesh.vertices.copy()
    faces_back = np.fliplr(faces_front) + len(vertices)
    vertices_double = np.vstack([vertices, vertices])
    faces_double = np.vstack([faces_front, faces_back])

    if hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None:
        colors = mesh.visual.vertex_colors.copy()
        colors_double = np.vstack([colors, colors])
    else:
        colors_double = None

    mesh_double = trimesh.Trimesh(
        vertices=vertices_double,
        faces=faces_double,
        vertex_colors=colors_double,
        process=False,
    )

    return mesh_double


def predict_from_motion(dx, dy, dz, rx, ry, rz, model_name):
    motion = np.array(
        [dx, dy, dz, rx / 180.0 * np.pi, ry / 180.0 * np.pi, rz / 180.0 * np.pi],
        dtype=np.float32,
    ).reshape(1, -1)
    assert motion.shape[1] == 6, "Motion must be 6D (x,y,z,rx,ry,rz)"

    outputs = ort_dict[model_name].run(None, {"motion": motion})
    force = np.array(outputs[0]).flatten()
    nodes = np.array(outputs[1]).reshape(-1, 3)

    df = pd.DataFrame(
        force.reshape(1, -1),
        columns=["Fx (N)", "Fy (N)", "Fz (N)", "Tx (Nmm)", "Ty (Nmm)", "Tz (Nmm)"],
    )
    df = df.style.format("{:.3f}")

    delta_vertices = np.zeros_like(vertices_dict[model_name])
    delta_vertices[def_index_dict[model_name] - 1] += nodes.reshape(-1, 3)
    new_vertices = vertices_dict[model_name] + delta_vertices

    rot_x_neg90 = np.array(
        [
            [1, 0, 0],
            [0, 0, 1],
            [0, -1, 0],
        ]
    )
    new_vertices = new_vertices @ rot_x_neg90.T

    norm_dist = np.clip(np.linalg.norm(delta_vertices, axis=1), cmin, cmax)
    colors = (colormap((norm_dist - cmin) / (cmax - cmin)) * 255).astype(np.uint8)

    mesh = trimesh.Trimesh(
        vertices=new_vertices,
        faces=faces_dict[model_name],
        vertex_colors=colors,
        process=False,
    )
    mesh = make_double_sided(mesh)

    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".glb")
    mesh.export(tmp_file.name, file_type="glb")

    return df, tmp_file.name


with gr.Blocks(title="FingerNet Demo") as demo:
    gr.Markdown("# 🖐️ FingerNet Demo\n\n")

    with gr.Row():
        with gr.Column():
            gr.Markdown("## Model Selection")
            model = gr.Dropdown(
                choices=list(model_dict.keys()),
                value=list(model_dict.keys())[0],
                label="Select Model",
                show_label=False,
            )

            gr.Markdown("## Motion Inputs")
            gr.Markdown("### Translation (mm)")
            dx = gr.Slider(-10, 10, 0, step=0.1, label="Dx")
            dy = gr.Slider(-10, 10, 0, step=0.1, label="Dy")
            dz = gr.Slider(-3, 3, 0, step=0.1, label="Dz")
            gr.Markdown("### Rotation (degrees)")
            rx = gr.Slider(-30, 30, 0, step=0.5, label="Rx")
            ry = gr.Slider(-30, 30, 0, step=0.5, label="Ry")
            rz = gr.Slider(-30, 30, 0, step=0.5, label="Rz")

        with gr.Column():
            gr.Markdown("## Force Outputs")
            output_force = gr.DataFrame(
                headers=[
                    "Fx (N)",
                    "Fy (N)",
                    "Fz (N)",
                    "Tx (Nmm)",
                    "Ty (Nmm)",
                    "Tz (Nmm)",
                ],
                label="Predicted Force",
                show_label=False,
                datatype=["number"] * 6,
                row_count=(1, "fixed"),
                interactive=False,
                column_widths=["16%"] * 6,
            )

            gr.Markdown("## Mesh Outputs")
            height = 500
            output_mesh = gr.Model3D(
                label="Mesh",
                show_label=False,
                height=height,
                camera_position=[45, 60, 150],
            )

    inputs = [dx, dy, dz, rx, ry, rz, model]
    for inp in inputs:
        inp.change(
            fn=predict_from_motion, inputs=inputs, outputs=[output_force, output_mesh]
        )

    demo.load(
        fn=predict_from_motion, inputs=inputs, outputs=[output_force, output_mesh]
    )

demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))