import numpy as np
import example_robot_data as erd
import pinocchio as pin
import meshcat
import meshcat.geometry as g
from pinocchio.visualize import MeshcatVisualizer
import time


def _actuated_dof_indices(model):
    idxs, labels, bounds = [], [], []
    for j in range(1, len(model.joints)):
        J = model.joints[j]
        if J.shortname() == "JointModelFreeFlyer":
            continue
        if J.nq == 1:
            i0 = J.idx_q
            lo = float(model.lowerPositionLimit[i0])
            hi = float(model.upperPositionLimit[i0])
            if not np.isfinite(lo): lo = -3.14
            if not np.isfinite(hi): hi =  3.14
            idxs.append(i0)
            labels.append(model.names[j])
            bounds.append((lo, hi))
    return idxs, labels, bounds

class LiveRobot:
    """
    Starts a local MeshCat ZMQ+Web server and keeps a persistent connection.
    The web UI is proxied at /meshcat/ by nginx (see nginx.conf).
    """
    def __init__(self, robot_name="talos"):
        self.vis = meshcat.Visualizer(zmq_url="tcp://127.0.0.1:6001")
        print(f"Zmq URL: {self.vis.window.zmq_url}, web URL: {self.vis.window.web_url}")

        self.rdata = erd.load(robot_name)
        self.model = self.rdata.model
        self.data = self.model.createData()
        
        self.vis.delete()
        time.sleep(0.5)  # wait for the robot to appear
        
        self.viz = MeshcatVisualizer(
            self.model, self.rdata.collision_model, self.rdata.visual_model
        )
        
        self.viz.initViewer(self.vis)        
        
        self.viz.clean()
        time.sleep(0.5)  # wait for the robot to appear
        
        self.viz.loadViewerModel(rootNodeName="robot")
        
        self.q = pin.neutral(self.model)
        try:
            self.viz.display(self.q)  
        except Exception as e:
            print(f"Error setting robot joints: {e}")

        
        self.idxs, self.labels, self.bounds = _actuated_dof_indices(self.model)


    def set_joints(self, vals):
        q = self.q.copy()
        for v, qi in zip(vals, self.idxs):
            q[qi] = float(v)
            
        try:
            self.viz.display(q)
            self.q = q
        except Exception as e:
            print(f"Error setting robot joints: {e}")
            

    def neutral(self):
        self.q = pin.neutral(self.model)
        self.viz.display(self.q)
        return [float(self.q[i]) for i in self.idxs]

    def iframe(self, width="100%", height=640):
        # nginx proxies the MeshCat web app at /meshcat/
        return f'<iframe src="/static/" style="width:{width};height:{height}px;border:0"></iframe>'