from dataclasses import dataclass, field
from typing import Optional, Tuple
import cadquery as Cq
from nhf import Role
from nhf.build import Model, target, assembly
import nhf.parts.springs as springs
from nhf.parts.joints import TorsionJoint
import nhf.utils

TOL = 1e-6

@dataclass
class ShoulderJoint(Model):

    shoulder_height: float = 100.0
    torsion_joint: TorsionJoint = field(default_factory=lambda: TorsionJoint(
        radius_track=18,
        radius_rider=18,
        groove_radius_outer=16,
        groove_radius_inner=13,
        track_disk_height=5.0,
        rider_disk_height=5.0,
        # M8 Axle
        radius_axle=3.0,
        # inner diameter = 9
        radius_spring=9/2 + 1.2,
        spring_thickness=1.3,
        spring_height=7.5,
    ))
    # Two holes on each side (top and bottom) are used to attach the shoulder
    # joint. This governs the distance between these two holes
    attach_dist: float = 25
    attach_diam: float = 8


    @target(name="shoulder-joint/parent")
    def parent(self,
               wing_root_wall_thickness: float = 5.0) -> Cq.Workplane:
        joint = self.torsion_joint
        # Thickness of the lip connecting this joint to the wing root
        lip_thickness = 10
        lip_width = 25
        lip_guard_ext = 40
        lip_guard_height = wing_root_wall_thickness + lip_thickness
        assert lip_guard_ext > joint.radius_track

        lip_guard = (
            Cq.Solid.makeBox(lip_guard_ext, lip_width, lip_guard_height)
            .located(Cq.Location((0, -lip_width/2 , 0)))
            .cut(Cq.Solid.makeCylinder(joint.radius_track, lip_guard_height))
        )
        result = (
            joint.track()
            .union(lip_guard, tol=1e-6)

            # Extrude the handle
            .copyWorkplane(Cq.Workplane(
                'YZ', origin=Cq.Vector((88, 0, wing_root_wall_thickness))))
            .rect(lip_width, lip_thickness, centered=(True, False))
            .extrude("next")

            # Connector holes on the lip
            .copyWorkplane(Cq.Workplane(
                'YX', origin=Cq.Vector((57, 0, wing_root_wall_thickness))))
            .hole(self.attach_diam)
            .moveTo(0, self.attach_dist)
            .hole(self.attach_diam)
        )
        result.moveTo(0, 0).tagPlane('conn0')
        result.moveTo(0, self.attach_dist).tagPlane('conn1')
        return result

    @property
    def child_height(self) -> float:
        """
        Calculates the y distance between two joint surfaces on the child side
        of the shoulder joint.
        """
        joint = self.torsion_joint
        return self.shoulder_height - 2 * joint.total_height + 2 * joint.rider_disk_height

    @target(name="shoulder-joint/child")
    def child(self,
              lip_height: float = 20.0,
              hole_dist: float = 10.0,
              spacer_hole_diam: float = 8.0) -> Cq.Assembly:
        """
        Creates the top/bottom shoulder child joint
        """

        joint = self.torsion_joint
        # Half of the height of the bridging cylinder
        dh = self.shoulder_height / 2 - joint.total_height
        core_start_angle = 30
        core_end_angle1 = 90
        core_end_angle2 = 180
        core_thickness = 2

        core_profile1 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, core_start_angle, core_end_angle1-core_start_angle)
            .segment((0, 0))
            .close()
            .assemble()
            .circle(joint.radius_rider - core_thickness, mode='s')
        )
        core_profile2 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, -core_start_angle, -(core_end_angle2-core_start_angle))
            .segment((0, 0))
            .close()
            .assemble()
            .circle(joint.radius_rider - core_thickness, mode='s')
        )
        core = (
            Cq.Workplane('XY')
            .placeSketch(core_profile1)
            .toPending()
            .extrude(dh * 2)
            .copyWorkplane(Cq.Workplane('XY'))
            .placeSketch(core_profile2)
            .toPending()
            .extrude(dh * 2)
            .translate(Cq.Vector(0, 0, -dh))
        )
        # Create the upper and lower lips
        lip_thickness = joint.rider_disk_height
        lip_ext = 40 + joint.radius_rider
        assert lip_height / 2 <= joint.radius_rider
        lip = (
            Cq.Workplane('XY')
            .box(lip_ext, lip_height, lip_thickness,
                 centered=(False, True, False))
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(radius=joint.radius_rider, height=lip_thickness,
                      centered=(True, True, False),
                      combine='cut')
            .faces(">Z")
            .workplane()
        )
        hole_x = lip_ext - hole_dist / 2
        for i in range(2):
            x = hole_x - i * hole_dist
            lip = lip.moveTo(x, 0).hole(spacer_hole_diam)
        for i in range(2):
            x = hole_x - i * hole_dist
            (
                lip
                .moveTo(x, 0)
                .tagPlane(f"conn{1 - i}")
            )

        loc_rotate = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        result = (
            Cq.Assembly()
            .add(core, name="core", loc=Cq.Location())
            .add(joint.rider(rider_slot_begin=-90, reverse_directrix_label=True), name="rider_top",
                 loc=Cq.Location((0, 0, dh), (0, 0, 1), -90))
            .add(joint.rider(rider_slot_begin=180), name="rider_bot",
                 loc=Cq.Location((0, 0, -dh), (0, 0, 1), -90) * loc_rotate)
            .add(lip, name="lip_top",
                 loc=Cq.Location((0, 0, dh)))
            .add(lip, name="lip_bot",
                 loc=Cq.Location((0, 0, -dh)) * loc_rotate)
        )
        return result

    @assembly()
    def assembly(self,
                 wing_root_wall_thickness: float = 5.0,
                 lip_height: float = 5.0,
                 hole_dist: float = 10.0,
                 spacer_hole_diam: float = 8.0
                 ) -> Cq.Assembly:
        directrix = 0
        result = (
            Cq.Assembly()
            .add(self.child(lip_height=lip_height,
                            hole_dist=hole_dist,
                            spacer_hole_diam=spacer_hole_diam),
                 name="child",
                 color=Role.CHILD.color)
            .constrain("child/core", "Fixed")
            .add(self.torsion_joint.spring(), name="spring_top",
                 color=Role.DAMPING.color)
            .add(self.parent(wing_root_wall_thickness),
                 name="parent_top",
                 color=Role.PARENT.color)
            .add(self.torsion_joint.spring(), name="spring_bot",
                 color=Role.DAMPING.color)
            .add(self.parent(wing_root_wall_thickness),
                 name="parent_bot",
                 color=Role.PARENT.color)
        )
        TorsionJoint.add_constraints(result,
                                     rider="child/rider_top",
                                     track="parent_top",
                                     spring="spring_top",
                                     directrix=directrix)
        TorsionJoint.add_constraints(result,
                                     rider="child/rider_bot",
                                     track="parent_bot",
                                     spring="spring_bot",
                                     directrix=directrix)
        return result.solve()


@dataclass
class Beam:
    """
    A I-shaped spine with two feet
    """

    foot_length: float = 40.0
    foot_width: float = 20.0
    foot_height: float = 5.0
    spine_thickness: float = 4.0
    spine_length: float = 10.0
    total_height: float = 50.0

    hole_diam: float = 8.0
    # distance between the centres of the two holes
    hole_dist: float = 24.0

    def __post_init__(self):
        assert self.spine_height > 0
        assert self.hole_diam + self.hole_dist < self.foot_length
        assert self.hole_dist - self.hole_diam >= self.spine_length

    @property
    def spine_height(self):
        return self.total_height - self.foot_height * 2

    def foot(self) -> Cq.Workplane:
        """
        just one foot
        """
        dx = self.hole_dist / 2
        result = (
            Cq.Workplane('XZ')
            .box(self.foot_length, self.foot_width, self.foot_height,
                 centered=(True, True, False))
            .faces(">Y")
            .workplane()
            .pushPoints([(dx, 0), (-dx, 0)])
            .hole(self.hole_diam)
        )
        plane = result.faces(">Y").workplane()
        plane.moveTo(dx, 0).tagPlane("conn1")
        plane.moveTo(-dx, 0).tagPlane("conn0")
        return result

    def beam(self) -> Cq.Assembly:
        beam = (
            Cq.Workplane('XZ')
            .box(self.spine_length, self.spine_thickness, self.spine_height)
        )
        h = self.spine_height / 2 + self.foot_height
        result = (
            Cq.Assembly()
            .add(beam, name="beam")
            .add(self.foot(), name="top",
                 loc=Cq.Location((0, h, 0)))
            .add(self.foot(), name="bot",
                 loc=Cq.Location((0, -h, 0), (1, 0, 0), 180))
        )
        return result


@dataclass
class DiskJoint(Model):
    """
    Sandwiched disk joint for the wrist and elbow
    """

    radius_housing: float = 22.0
    radius_disk: float = 20.0
    radius_spring: float = 9 / 2
    radius_axle: float = 3.0

    housing_thickness: float = 5.0
    disk_thickness: float = 5.0
    # Gap between disk and the housing
    #disk_thickness_gap: float = 0.1
    spring_thickness: float = 1.3
    spring_height: float = 6.5
    spring_tail_length: float = 45.0

    # Spring angle at 0 degrees of movement
    spring_angle: float = 30.0
    # Angle at which the spring exerts no torque
    spring_angle_neutral: float = 90.0
    spring_angle_shift: float = 30
    wall_inset: float = 2.0

    # Angular span of movement
    movement_angle: float = 120.0
    # Angular span of tongue on disk
    tongue_span: float = 30.0
    tongue_length: float = 10.0

    generate_inner_wall: bool = False


    def __post_init__(self):
        super().__init__(name="disk-joint")
        assert self.housing_thickness > self.wall_inset
        assert self.radius_housing > self.radius_disk
        assert self.radius_disk > self.radius_axle
        assert self.housing_upper_carve_offset > 0

    def spring(self):
        return springs.torsion_spring(
            radius=self.radius_spring,
            height=self.spring_height,
            thickness=self.spring_thickness,
            tail_length=self.spring_tail_length,
            right_handed=False,
        )

    @property
    def neutral_movement_angle(self) -> Optional[float]:
        a = self.spring_angle_neutral - self.spring_angle
        if 0 <= a and a <= self.movement_angle:
            return a
        return None

    @property
    def total_thickness(self) -> float:
        return self.housing_thickness * 2 + self.disk_thickness

    @property
    def opening_span(self) -> float:
        return self.movement_angle + self.tongue_span

    @property
    def housing_upper_carve_offset(self) -> float:
        """
        Distance between the spring track and the outside of the upper housing
        """
        return self.housing_thickness + self.disk_thickness - self.spring_height

    @property
    def housing_upper_dz(self) -> float:
        """
        Distance between the default upper housing location and the median line
        """
        return self.total_thickness / 2 - self.housing_thickness

    @property
    def radius_spring_internal(self):
        return self.radius_spring - self.spring_thickness

    @target(name="disk")
    def disk(self) -> Cq.Workplane:
        cut = (
            Cq.Solid.makeBox(
                length=self.spring_tail_length,
                width=self.spring_thickness,
                height=self.disk_thickness,
            )
            .located(Cq.Location((0, self.radius_spring_internal, 0)))
            .rotate((0, 0, 0), (0, 0, 1), self.spring_angle_shift)
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
                centered=(True, True, False)
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_spring,
                centered=(True, True, False),
                combine='cut',
            )
            .cut(cut)
        )
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        plane.tagPlane("dir", direction="+X")
        plane.workplane(offset=self.disk_thickness).tagPlane("mate_top")
        result.copyWorkplane(Cq.Workplane('YX')).tagPlane("mate_bot")

        radius_tongue = self.radius_disk + self.tongue_length
        tongue = (
            Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=radius_tongue,
                angleDegrees=self.tongue_span,
            ).cut(Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
            ))
        )
        result = result.union(tongue, tol=TOL)
        return result

    def wall(self) -> Cq.Compound:
        height = self.disk_thickness + self.wall_inset
        wall = Cq.Solid.makeCylinder(
            radius=self.radius_housing,
            height=height,
            angleDegrees=360 - self.opening_span,
        ).cut(Cq.Solid.makeCylinder(
            radius=self.radius_disk,
            height=height,
        )).rotate((0, 0, 0), (0, 0, 1), self.opening_span)
        return wall

    @target(name="housing-lower")
    def housing_lower(self) -> Cq.Workplane:
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
            .cut(Cq.Solid.makeCylinder(
                radius=self.radius_axle,
                height=self.housing_thickness,
            ))
        )
        result.faces(">Z").tag("mate")
        result.faces(">Z").workplane().tagPlane("dirX", direction="+X")
        # two directional vectors are required to make the angle constrain
        # unambiguous
        result.faces(">Z").workplane().tagPlane("dirY", direction="+Y")
        result = result.cut(
            self
            .wall()
            .located(Cq.Location((0, 0, self.disk_thickness - self.wall_inset)))
            #.rotate((0, 0, 0), (1, 0, 0), 180)
            #.located(Cq.Location((0, 0, self.disk_thickness + self.housing_thickness)))
        )
        return result

    @target(name="housing-upper")
    def housing_upper(self) -> Cq.Workplane:
        carve = (
            Cq.Solid.makeCylinder(
                radius=self.radius_spring,
                height=self.housing_thickness
            ).fuse(Cq.Solid.makeBox(
                length=self.spring_tail_length,
                width=self.spring_thickness,
                height=self.housing_thickness
            ).located(Cq.Location((0, -self.radius_spring, 0))))
        ).rotate((0, 0, 0), (0, 0, 1), self.spring_angle - self.spring_angle_shift)
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
        )
        result.faces("<Z").tag("mate")
        result.copyWorkplane(Cq.Workplane('XY')).tagPlane("dir", direction="+X")
        result = result.faces(">Z").hole(self.radius_axle * 2)

        # tube which holds the spring interior
        if self.generate_inner_wall:
            tube = (
                Cq.Solid.makeCylinder(
                    radius=self.radius_spring_internal,
                    height=self.disk_thickness + self.housing_thickness,
                ).cut(Cq.Solid.makeCylinder(
                    radius=self.radius_axle,
                    height=self.disk_thickness + self.housing_thickness,
                ))
            )
            result = result.union(tube)
        wall = (
            self.wall()
            .located(Cq.Location((0, 0, -self.disk_thickness-self.wall_inset)))
        )
        result = (
            result
            .cut(carve.located(Cq.Location((0, 0, -self.housing_upper_carve_offset))))
            .union(wall, tol=TOL)
        )
        return result

    def add_constraints(self,
                        assembly: Cq.Assembly,
                        housing_lower: str,
                        housing_upper: str,
                        disk: str,
                        angle: float,
                        ) -> Cq.Assembly:
        return (
            assembly
            .constrain(f"{disk}?mate_bot", f"{housing_lower}?mate", "Plane")
            .constrain(f"{disk}?mate_top", f"{housing_upper}?mate", "Plane")
            .constrain(f"{housing_lower}?dirX", f"{housing_upper}?dir", "Axis", param=0)
            .constrain(f"{housing_lower}?dirX", f"{disk}?dir", "Axis", param=angle)
            .constrain(f"{housing_lower}?dirY", f"{disk}?dir", "Axis", param=angle - 90)
        )


    def assembly(self, angle: Optional[float] = 0) -> Cq.Assembly:
        if angle is None:
            angle = self.movement_angle
            if angle is None:
                angle = 0
        else:
            assert 0 <= angle <= self.movement_angle
        result = (
            Cq.Assembly()
            .add(self.disk(), name="disk", color=Role.CHILD.color)
            .add(self.housing_lower(), name="housing_lower", color=Role.PARENT.color)
            .add(self.housing_upper(), name="housing_upper", color=Role.CASING.color)
            #.constrain("housing_lower", "Fixed")
        )
        result = self.add_constraints(
            result,
            housing_lower="housing_lower",
            housing_upper="housing_upper",
            disk="disk",
            angle=angle,
        )
        return result.solve()


@dataclass
class ElbowJoint:
    """
    Creates the elbow and wrist joints.

    This consists of a disk joint, where each side of the joint has mounting
    holes for connection to the exoskeleton. Each side 2 mounting feet on the
    top and bottom, and each foot has 2 holes.

    On the parent side, additional bolts are needed to clamp the two sides of
    the housing together.
    """

    disk_joint: DiskJoint = field(default_factory=lambda: DiskJoint(
        movement_angle=60,
    ))

    # Distance between the child/parent arm to the centre
    child_arm_radius: float = 40.0
    parent_arm_radius: float = 40.0

    child_beam: Beam = field(default_factory=lambda: Beam())
    parent_beam: Beam = field(default_factory=lambda: Beam(
        spine_thickness=8.0,
    ))
    parent_arm_span: float = 40.0
    # Angle of the beginning of the parent arm
    parent_arm_angle: float = 180.0
    parent_binding_hole_radius: float = 30.0

    # Size of the mounting holes
    hole_diam: float = 8.0

    def __post_init__(self):
        assert self.child_arm_radius > self.disk_joint.radius_housing
        assert self.parent_arm_radius > self.disk_joint.radius_housing
        self.disk_joint.tongue_length = self.child_arm_radius - self.disk_joint.radius_disk
        assert self.disk_joint.movement_angle < self.parent_arm_angle < 360 - self.parent_arm_span
        assert self.parent_binding_hole_radius - self.hole_diam / 2 > self.disk_joint.radius_housing

    def child_joint(self) -> Cq.Assembly:
        angle = -self.disk_joint.tongue_span / 2
        dz = self.disk_joint.disk_thickness / 2
        # We need to ensure the disk is on the "other" side so
        flip = Cq.Location((0, 0, 0), (0, 0, 1), 180)
        result = (
            self.child_beam.beam()
            .add(self.disk_joint.disk(), name="disk",
                 loc=flip * Cq.Location((-self.child_arm_radius, 0, -dz), (0, 0, 1), angle))
            #.constrain("disk", "Fixed")
            #.constrain("top", "Fixed")
            #.constrain("bot", "Fixed")
            #.solve()
        )
        return result

    def parent_joint_lower(self) -> Cq.Workplane:
        return self.disk_joint.housing_lower()

    def parent_joint_upper(self):
        axial_offset = Cq.Location((self.parent_arm_radius, 0, 0))
        housing_dz = self.disk_joint.housing_upper_dz
        conn_h = self.parent_beam.spine_thickness
        connector = (
            Cq.Solid.makeCylinder(
                height=conn_h,
                radius=self.parent_arm_radius,
                angleDegrees=self.parent_arm_span)
            .cut(Cq.Solid.makeCylinder(
                height=conn_h,
                radius=self.disk_joint.radius_housing,
            ))
            .located(Cq.Location((0, 0, -conn_h / 2)))
            .rotate((0,0,0), (0,0,1), 180-self.parent_arm_span / 2)
        )
        housing = self.disk_joint.housing_upper()
        result = (
            self.parent_beam.beam()
            .add(housing, name="housing",
                 loc=axial_offset * Cq.Location((0, 0, housing_dz)))
            .add(connector, name="connector",
                 loc=axial_offset)
            #.constrain("housing", "Fixed")
            #.constrain("connector", "Fixed")
            #.solve()
        )
        return result

    def assembly(self, angle: float = 0) -> Cq.Assembly:
        result = (
            Cq.Assembly()
            .add(self.child_joint(), name="child", color=Role.CHILD.color)
            .add(self.parent_joint_lower(), name="parent_lower", color=Role.CASING.color)
            .add(self.parent_joint_upper(), name="parent_upper", color=Role.PARENT.color)
            #.constrain("child/disk?mate_bot", "Fixed")
        )
        result = self.disk_joint.add_constraints(
            result,
            housing_lower="parent_lower",
            housing_upper="parent_upper/housing",
            disk="child/disk",
            angle=angle,
        )
        return result.solve()

if __name__ == '__main__':
    p = ShoulderJoint()
    p.build_all()
    p = DiskJoint()
    p.build_all()