from time import strftime
from lxml import etree
from ...components.real.biomechanical_model_real import BiomechanicalModelReal
from ...utils.enums import Translations, Rotations
from ..abstract_model_writer import AbstractModelWriter


class OpensimModelWriter(AbstractModelWriter):

    def write(self, model: "BiomechanicalModelReal") -> None:
        """
        Write the OpenSim model to file

        Parameters
        ----------
        model
            The biomechanical model to write
        """
        self.check_that_model_is_osimable(model)

        root = etree.Element("OpenSimDocument", Version="40000")
        model_elem = etree.SubElement(root, "Model", name="model")

        self._write_credits(model_elem, model)
        self._write_units(model_elem)
        self._write_gravity(model_elem, model)
        self._write_ground(model_elem, model)
        self._write_body_set(model_elem, model)
        self._write_joint_set(model_elem, model)
        self._write_marker_set(model_elem, model)
        self._write_force_set(model_elem, model)

        tree = etree.ElementTree(root)
        tree.write(self.filepath, pretty_print=True, xml_declaration=True, encoding="UTF-8")

    def check_that_model_is_osimable(self, model: BiomechanicalModelReal) -> None:
        """Check that the model can be written to OpenSim format"""

        if model.nb_contacts > 0:
            raise NotImplementedError("Writing models with contacts to OpenSim format is not yet implemented.")

        if model.nb_imus > 0:
            raise NotImplementedError("Writing models with IMUs to OpenSim format is not yet implemented.")

        for segment in model.segments:
            if segment.mesh is not None and not self.with_mesh:
                raise NotImplementedError(
                    f"Writing models with segment meshes to OpenSim format is not yet implemented (segment: {segment.name})."
                )

    def _write_credits(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write credits and publication information"""
        credits = etree.SubElement(model_elem, "credits")
        credits.text = model.header + f"Model generated by BioBuddy on {strftime('%Y-%m-%d %H:%M:%S')}"

    def _write_units(self, model_elem: etree.Element):
        """Write unit information"""
        length_units = etree.SubElement(model_elem, "length_units")
        length_units.text = "meters"

        force_units = etree.SubElement(model_elem, "force_units")
        force_units.text = "N"

    def _write_gravity(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write gravity vector"""
        gravity_elem = etree.SubElement(model_elem, "gravity")
        if model.gravity is not None:
            gravity_elem.text = f"{model.gravity[0, 0]:.8f} {model.gravity[1, 0]:.8f} {model.gravity[2, 0]:.8f}"
        else:
            gravity_elem.text = "0 -9.80665 0"

    def _write_ground(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """
        Write ground body.
        Note that we assume that the ground segment is named "ground" in the model.
        """
        ground_elem = etree.SubElement(model_elem, "Ground", name="ground")
        ground_geometry = etree.SubElement(ground_elem, "FrameGeometry", name="frame_geometry")
        socket_frame = etree.SubElement(ground_geometry, "socket_frame")
        socket_frame.text = ".."
        if "ground" in model.segment_names:
            ground_segment = model.segments["ground"]
            attached_geometry = etree.SubElement(ground_geometry, "attached_geometry")
            socket_frame = etree.SubElement(attached_geometry, "socket_frame")
            socket_frame.text = ".."
            if ground_segment.mesh_file is not None:
                scale_factors = etree.SubElement(attached_geometry, "scale_factors")
                ground_scale_factors = ground_segment.mesh_file.mesh_scale
                scale_factors.text = (
                    f"{ground_scale_factors[0]:.8f} {ground_scale_factors[1]:.8f} {ground_scale_factors[2]:.8f}"
                )

            if ground_segment.nb_markers > 0:
                for marker in ground_segment.markers:
                    marker_elem = marker.to_osim()
                    attached_geometry.append(marker_elem)

    def _write_body_set(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write all body segments"""
        body_set = etree.SubElement(model_elem, "BodySet", name="bodyset")
        objects = etree.SubElement(body_set, "objects")

        for segment in model.segments:
            if segment.name == "ground" or segment.name == "root":
                continue

            body_elem = segment.to_osim(with_mesh=self.with_mesh)
            objects.append(body_elem)

    def _write_joint_set(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write all joints"""
        joint_set = etree.SubElement(model_elem, "JointSet", name="jointset")
        objects = etree.SubElement(joint_set, "objects")

        for segment in model.segments:
            if segment.name == "ground" or segment.name == "root":
                continue

            parent_name = segment.parent_name  # if segment.parent_name != "base" else "ground"
            if segment.nb_q > 0:
                joint_elem = self._create_custom_joint(segment, parent_name)
            else:
                joint_elem = self._create_weld_joint(segment, parent_name)

            objects.append(joint_elem)

    def _create_weld_joint(self, segment, parent_name: str):
        """Create a WeldJoint (no DOFs)"""
        joint = etree.Element("WeldJoint", name=f"{segment.name}_joint")

        socket_parent = etree.SubElement(joint, "socket_parent_frame")
        if parent_name == "ground":
            socket_parent.text = "../ground"
        else:
            socket_parent.text = f"bodyset/{parent_name}"

        socket_child = etree.SubElement(joint, "socket_child_frame")
        socket_child.text = f"bodyset/{segment.name}"

        frames = etree.SubElement(joint, "frames")

        # Parent offset frame
        parent_frame = etree.SubElement(frames, "PhysicalOffsetFrame", name=f"{parent_name}")
        parent_socket = etree.SubElement(parent_frame, "socket_parent")
        parent_socket.text = f"bodyset/{parent_name}"

        # Use the segment coordinate system's to_osim method
        translation_data, _ = segment.segment_coordinate_system.to_osim()

        translation = etree.SubElement(parent_frame, "translation")
        translation.text = f"{translation_data[0]:.8f} {translation_data[1]:.8f} {translation_data[2]:.8f}"

        orientation = etree.SubElement(parent_frame, "orientation")
        orientation.text = f"0 0 0"

        # Child offset frame
        child_frame = etree.SubElement(frames, "PhysicalOffsetFrame", name=f"{segment.name}")
        child_socket = etree.SubElement(child_frame, "socket_parent")
        child_socket.text = f"bodyset/{segment.name}"

        child_translation = etree.SubElement(child_frame, "translation")
        child_translation.text = "0 0 0"

        child_orientation = etree.SubElement(child_frame, "orientation")
        child_orientation.text = "0 0 0"

        return joint

    def _create_custom_joint(self, segment, parent_name: str):
        """Create a CustomJoint with DOFs"""
        joint = etree.Element("CustomJoint", name=f"{segment.name}_joint")

        socket_parent = etree.SubElement(joint, "socket_parent_frame")
        if parent_name == "ground":
            socket_parent.text = "../ground"
        else:
            socket_parent.text = f"bodyset/{parent_name}"

        socket_child = etree.SubElement(joint, "socket_child_frame")
        socket_child.text = f"bodyset/{segment.name}"

        frames = etree.SubElement(joint, "frames")

        # Parent offset frame
        parent_frame = etree.SubElement(frames, "PhysicalOffsetFrame", name=f"{parent_name}")
        parent_socket = etree.SubElement(parent_frame, "socket_parent")
        parent_socket.text = f"bodyset/{parent_name}"

        # Segment coordinate system
        translation_data, _ = segment.segment_coordinate_system.to_osim()

        translation = etree.SubElement(parent_frame, "translation")
        translation.text = f"{translation_data[0]:.8f} {translation_data[1]:.8f} {translation_data[2]:.8f}"

        orientation = etree.SubElement(parent_frame, "orientation")
        orientation.text = f"0 0 0"

        # Child offset frame
        child_frame = etree.SubElement(frames, "PhysicalOffsetFrame", name=f"{segment.name}")
        child_socket = etree.SubElement(child_frame, "socket_parent")
        child_socket.text = f"bodyset/{segment.name}"

        child_translation = etree.SubElement(child_frame, "translation")
        child_translation.text = "0 0 0"

        child_orientation = etree.SubElement(child_frame, "orientation")
        child_orientation.text = "0 0 0"

        # Initialize coordinates
        coordinates = etree.SubElement(joint, "coordinates")

        # Initialize spatial transform
        spatial_transform = etree.SubElement(joint, "SpatialTransform")

        dof_counter = 0

        # Rotations (first 3 transform axes)
        for i, axis_name in enumerate(["X", "Y", "Z"]):
            transform_axis = etree.SubElement(spatial_transform, "TransformAxis", name=f"rotation{i+1}")

            axis_elem = etree.SubElement(transform_axis, "axis")
            axis_vector = [0, 0, 0]
            axis_vector[i] = 1
            axis_elem.text = f"{axis_vector[0]} {axis_vector[1]} {axis_vector[2]}"

            if segment.rotations != Rotations.NONE and axis_name.lower() in segment.rotations.value:
                coord_name = segment.dof_names[dof_counter]
                dof_counter += 1

                coord_transform_axis = etree.SubElement(transform_axis, "coordinates")
                coord_transform_axis.text = coord_name
                coord_elem = etree.SubElement(coordinates, "Coordinate", name=coord_name)

                default_value = etree.SubElement(coord_elem, "default_value")
                default_value.text = "0"

                default_speed_value = etree.SubElement(coord_elem, "default_speed_value")
                default_speed_value.text = "0"

                locked = etree.SubElement(coord_elem, "locked")
                locked.text = "false"

                prescribed = etree.SubElement(coord_elem, "prescribed")
                prescribed.text = "false"

                if segment.q_ranges is not None:
                    idx = segment.dof_names.index(coord_name)
                    range_elem = etree.SubElement(coord_elem, "range")
                    range_elem.text = f"{segment.q_ranges.min_bound[idx]:.8f} {segment.q_ranges.max_bound[idx]:.8f}"

                    clamped = etree.SubElement(coord_elem, "clamped")
                    clamped.text = "true"

                locked = etree.SubElement(coord_elem, "locked")
                locked.text = "false"
            else:
                function = etree.SubElement(transform_axis, "function")
                constant = etree.SubElement(function, "Constant")
                value = etree.SubElement(constant, "value")
                value.text = "0"

        # Translations (last 3 transform axes)
        for i, axis_name in enumerate(["X", "Y", "Z"]):
            transform_axis = etree.SubElement(spatial_transform, "TransformAxis", name=f"translation{i+1}")

            axis_elem = etree.SubElement(transform_axis, "axis")
            axis_vector = [0, 0, 0]
            axis_vector[i] = 1
            axis_elem.text = f"{axis_vector[0]} {axis_vector[1]} {axis_vector[2]}"

            if segment.translations != Translations.NONE and axis_name.lower() in segment.translations.value:
                coord_name = segment.dof_names[dof_counter]
                dof_counter += 1

                coord_transform_axis = etree.SubElement(transform_axis, "coordinates")
                coord_transform_axis.text = coord_name
                coord_elem = etree.SubElement(coordinates, "Coordinate", name=coord_name)

                default_value = etree.SubElement(coord_elem, "default_value")
                default_value.text = "0"

                if segment.q_ranges is not None:
                    idx = segment.dof_names.index(coord_name)
                    range_elem = etree.SubElement(coord_elem, "range")
                    range_elem.text = f"{segment.q_ranges.min_bound[idx]:.8f} {segment.q_ranges.max_bound[idx]:.8f}"

                    clamped = etree.SubElement(coord_elem, "clamped")
                    clamped.text = "false"

                locked = etree.SubElement(coord_elem, "locked")
                locked.text = "false"
            else:
                function = etree.SubElement(transform_axis, "function")
                constant = etree.SubElement(function, "Constant")
                value = etree.SubElement(constant, "value")
                value.text = "0"

        return joint

    def _write_marker_set(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write all markers"""
        marker_set = etree.SubElement(model_elem, "MarkerSet")
        objects = etree.SubElement(marker_set, "objects")

        for segment in model.segments:
            if segment.name == "ground":
                continue

            for marker in segment.markers:
                marker_elem = marker.to_osim()
                objects.append(marker_elem)

    def _write_force_set(self, model_elem: etree.Element, model: BiomechanicalModelReal):
        """Write all muscles"""
        force_set = etree.SubElement(model_elem, "ForceSet", name="forceset")
        objects = etree.SubElement(force_set, "objects")

        for muscle_group in model.muscle_groups:
            for muscle in muscle_group.muscles:
                muscle_elem = muscle.to_osim()
                objects.append(muscle_elem)
