from typing import List, Tuple

import numpy as np
from pathlib import Path
from pydrake.geometry import Meshcat, MeshcatVisualizer, MeshcatVisualizerParams
from pydrake.geometry import Role as DrakeRole
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph, MultibodyPlant
from pydrake.systems.framework import DiagramBuilder, Diagram, Context

from pydrake.systems.primitives import ConstantVectorSource, VectorLogSink

# Internal Imports
from brom_drake.robots import find_base_link_name_in
from brom_drake.productions.types import BaseProduction
from brom_drake.productions import ProductionID
from brom_drake.productions.roles.role import Role
from brom_drake.utils import Performer, AddMultibodyTriad
from .show_me_system import ShowMeSystem

class ShowMeThisModel(BaseProduction):
    """
    **Description**
    
    A production which allows the user to visualize a specified model
    in Meshcat, with the ability to set desired joint positions, if the model is
    a robot/articulated.
    
    **Parameters**
    
    path_to_model: str
        The path to the model file to be visualized.
        TODO: Add support for Path type. Make sure to add tests!!

    with_these_joint_positions: List[float], optional
        A list of joint positions to set the model to when visualized.
        The length of this list must match the number of positions in the model.
        If None, the model will be visualized in its default configuration.
        Default is None.

    base_link_name: str, optional
        The name of the base link of the model. If None, the production will
        attempt to automatically determine the base link name.
        Default is None.

    time_step: float, optional
        The time step to use for the MultibodyPlant.
        Default is 1e-3.

    meshcat_port_number: int, optional
        The port number to use for Meshcat visualization.
        If None, Meshcat visualization will be disabled.
        Default is 7001.

    show_collision_geometries: bool, optional
        A boolean flag indicating whether to show collision geometries in Meshcat
        INSTEAD of the visual geometries.
        Default is False.

    **Usage**

    The production can be used to show both simple models (e.g., single objects)
    or articulated models (e.g., robots). 
    
    To show a simple model, simply provide the path to the model file. ::

        from brom_drake.productions import ShowMeThisModel

        # Some stuff here

        urdf_file_path = "path/to/your/model.urdf"

        # Create the production
        production = ShowMeThisModel(
            urdf_file_path,
            time_step=1e-3,
        )

        # Add the cast to the production, build the diagram, and simulate
        diagram, diagram_context = production.add_cast_and_build()

        # Set up simulation
        simulator = Simulator(diagram, diagram_context)
        simulator.set_target_realtime_rate(1.0)
        simulator.set_publish_every_time_step(False)

    If the model is articulated, then you can also provide desired joint positions. ::

        from brom_drake.productions import ShowMeThisModel

        # Some stuff here

        urdf_file_path = "path/to/your/articulated_model.urdf"
        desired_joint_positions = [0.0, 1.0, -1.0, 0.5]  # Example joint positions

        # Create the production
        production = ShowMeThisModel(
            urdf_file_path,
            with_these_joint_positions=desired_joint_positions,
            time_step=1e-3,
        )

        # Add the cast to the production, build the diagram, and simulate
        diagram, diagram_context = production.add_cast_and_build()

        # Set up simulation
        simulator = Simulator(diagram, diagram_context)
        simulator.set_target_realtime_rate(1.0)
        simulator.set_publish_every_time_step(False)

    """
    def __init__(
        self,
        path_to_model: str|Path,
        with_these_joint_positions: List[float] = None,
        base_link_name: str = None,
        time_step: float = 1e-3,
        meshcat_port_number: int = 7001, # Usually turn off for CI (i.e., make it None)
        show_collision_geometries: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        # Copy all parameters into the production
        if type(path_to_model) is str:
            path_to_model = Path(path_to_model)

        self.path_to_model = path_to_model
        self.q_des = with_these_joint_positions
        self.base_link_name = base_link_name
        self.time_step = time_step
        self.meshcat_port_number = meshcat_port_number
        self.show_collision_geometries = show_collision_geometries

        # Check that the model exists
        if not self.path_to_model.exists():
            raise ValueError(f"The path \"{self.path_to_model}\" does not exist! Can not show a model that does not exist!")

        # If the base link name is not provided,
        # then we will try to do a smart search for it.
        if self.base_link_name is None:
            self.base_link_name = find_base_link_name_in(self.path_to_model)

        # Add Plant and Scene Graph for easy simulation
        self.plant, self.scene_graph = AddMultibodyPlantSceneGraph(
            self.builder,
            time_step=self.time_step,
        )
        self.plant.set_name("ShowMeThisModel_plant")

        self.meshcat = None
        self.model_index, self.model_name = None, None
        self.show_me_system = None

    def add_supporting_cast(self):
        """
        **Description**
        
        This method updates the production's
        - MultibodyPlant with the user's model
        - Using the ShowMeSystem to freeze the model in place
        - Connects to Meshcat, if requested

        This is required to be implemented by all children of the BaseProduction class.
        """
        # Setup

        # Add Model
        model_idcs = Parser(plant=self.plant).AddModels(
            str(self.path_to_model)
        )
        assert len(model_idcs) == 1, \
            f"Only one model should be added in this production;" + \
            f" received {len(model_idcs)} in the file {self.path_to_model}."
        self.model_index = model_idcs[0]
        self.model_name = self.plant.GetModelInstanceName(model_idcs[0])

        # Collect the expected number of actuated joints
        n_dofs = self.find_number_of_positions_in_model()
        if self.q_des is None:
            self.q_des = np.zeros(n_dofs)
        else:
            assert len(self.q_des) == n_dofs, \
                f"Expected {n_dofs} joint positions; received {len(self.q_des)}."

        # Add ShowMeSystem
        self.show_me_system = self.builder.AddSystem(
            ShowMeSystem(
                plant=self.plant,
                model_index=self.model_index,
                desired_joint_positions=np.array(self.q_des),
            ),
        )

        self.create_and_connect_source_of_joint_positions()

        # Add Sink
        output_joints_sink = self.builder.AddSystem(
            VectorLogSink(len(self.q_des)),
        )

        # Connect the system to the sink
        self.builder.Connect(
            self.show_me_system.get_output_port(0),
            output_joints_sink.get_input_port(0),
        )

        self.add_multibody_triads()

        # Weld the base link to the world frame
        self.plant.WeldFrames(
            self.plant.world_frame(),
            self.plant.GetFrameByName(self.base_link_name, self.model_index),
        )

        # Connect to Meshcat, if requested
        if self.meshcat_port_number is not None:
            self.meshcat = Meshcat(port=self.meshcat_port_number)  # Object provides an interface to Meshcat
            m_visualizer = MeshcatVisualizer(
                self.meshcat,
            )
            params = MeshcatVisualizerParams(
                role=DrakeRole.kIllustration,
            )
            if self.show_collision_geometries:
                params = MeshcatVisualizerParams(
                    role=DrakeRole.kProximity,
                )

            m_visualizer.AddToBuilder(
                self.builder, self.scene_graph, self.meshcat,
                params=params,
            )

        # Finalize plant and connect it to system
        self.plant.Finalize()

    def add_cast_and_build(
        self,
        cast: Tuple[Role, Performer] = [],
    ) -> Tuple[Diagram, Context]:
        """
        **Description**
        
        This method builds the production's diagram and context,
        and assigns the plant context to the internal ShowMeSystem.
        
        **Parameters**
        
        cast: Tuple[Role, Performer], optional
            The cast to be added to the production.
            Default is an empty tuple.
            
        **Returns**

        diagram: Diagram
            The built Diagram of the production.

        diagram_context: Context
            The built Context of the production.
        """
        super().add_cast_and_build(cast)

        # Assign the diagram context to the internal show_me_system
        self.show_me_system.mutable_plant_context = self.plant.GetMyMutableContextFromRoot(
            self.diagram_context,
        )

        return self.diagram, self.diagram_context

    def add_multibody_triads(self):
        """
        **Description**
        
        This method will add triads to the world frame
        and any other relevant frames.
        """
        # Setup

        # Input Checking
        assert self.model_index is not None, \
            "The model index is not yet set. Please call add_supporting_cast() first."

        # Add A Triad to base?
        AddMultibodyTriad(self.plant.world_frame(), self.scene_graph)

    def create_and_connect_source_of_joint_positions(self):
        """
        **Description**
        
        This method will create and connect a source of joint positions
        to the ShowMeSystem.

        **Assumptions**

        - ``self.show_me_system`` is not None
        - ``self.q_des`` is not None
        """
        # Setup

        assert self.show_me_system is not None, \
            "The show_me_system is not yet created. Please call add_supporting_cast() first."

        assert self.q_des is not None, \
            "The desired joint positions are not yet set. Please provide them at initialization."

        # Add Source
        desired_joint_positions_source = self.builder.AddSystem(
            ConstantVectorSource(np.array(self.q_des)),
        )

        # Connect the source to the system
        self.builder.Connect(
            desired_joint_positions_source.get_output_port(0),
            self.show_me_system.get_input_port(0),
        )

    def find_number_of_positions_in_model(
        self,
    ) -> int:
        """
        **Description**
        
        This method will return the number of positions in the model
        as defined by the user (through the path_to_model).

        **Returns**
        
        n_positions: int
            The number of positions in the model.
        """
        # Setup

        # Create a shadow plant
        shadow_plant = MultibodyPlant(self.time_step)
        model_idcs = Parser(plant=shadow_plant).AddModels(
            str(self.path_to_model)
        )
        shadow_model_idx = model_idcs[0]

        # Weld the base link to the world frame
        shadow_plant.WeldFrames(
            shadow_plant.world_frame(),
            shadow_plant.GetFrameByName(self.base_link_name, shadow_model_idx),
        )

        # Finalize the shadow plant
        shadow_plant.Finalize()

        return shadow_plant.num_positions()

    @property
    def id(self) -> ProductionID:
        """
        **Description**
        
        This property returns the unique identifier for this production.
        (i.e., ProductionID.kShowMeThisModel)

        This is required to be implemented by all children of the BaseProduction class.

        **Returns**
        
        id: ProductionID
            The unique identifier for this production.
        """
        return ProductionID.kShowMeThisModel

    @property
    def suggested_roles(self) -> List[Role]:
        """
        **Description**

        This property returns an empty list, because there are no suggested
        roles for this production.

        This is required to be implemented by all children of the BaseProduction class.

        **Returns**

        roles: List[Role]
            A list of suggested roles for this production.
        """
        return []
